FedLite: A Scalable Approach for Federated Learning on Resource-constrained Clients
Authors: Jianyu Wang
Hang Qi
Ankit Singh Rawat
Sashank J. Reddi
Sagar M. Waghmare
Felix Yu
Gauri Joshi
Venue: ICLR 2022 Submitted (2022)

Abstract: In classical federated learning, the clients contribute to the overall training by communicating local updates for the underlying model on their private data to a coordinating server.  However, updating and communicating the entire model becomes prohibitively expensive when resource-constrained clients collectively aim to train a large machine learning model. Split learning provides a natural solution in such a setting, where only a (small) part of the model is stored and trained on clients while the remaining (large) part of the model only stays at the servers. Unfortunately, the model partitioning employed in split learning significantly increases the communication cost compared to the classical federated learning algorithms. This paper addresses this issue by proposing an end-to-end training framework that relies on a novel vector quantization scheme accompanied by a gradient correction method to reduce the additional communication cost associated with split learning. An extensive empirical evaluation on standard image and text benchmarks shows that the proposed method can achieve up to $490\times$ communication cost reduction with minimal drop in accuracy, and enables a desirable performance vs. communication trade-off. 

Under review as a conference paper at ICLR 2022
FEDLITE: A SCALABLE APPROACH FOR FEDERATED LEARNING ON RESOURCE-COOSTRRIIND) CLIENTS
Anonymous authors Paper under double-blind review
ABSTRACT
In classical federated learning, the clients contribute to the overall training by communicating local updates for the underlying model on their private data to coordinating server. However, updating and communicating the entire model becomes prohibitively expensive when resource-constraine
clients collectively aim to train large machine learning model. Split learning provides natural solution in such a setting, where only a (small) part of the model is stored and trained on clients while the remaining (large) part of the model only stays the servers. Unfortunately, the model partitioning employed in split learning significantly increases the communication cost compared to the classical federated learning algorithms. This paper addresses this issue by proposing an end-to-end training framework that relies on novel vector quantization scheme accompanied by gradient correction method to reduce the additional communication cost associated with split learning. An extensive empirical evaluation on standard image and text benchmarks shows that the proposed method can achieve up to x90 communication cost reduction with minimal drop in accuracy, and enables a desirable performance communication trade-off.
INTRODUCTION
Federated learning (FL) is an emerging field that collaboratively trains machine learning models on decentralized data (Li et al., 2019; Kairouz et al.. 2019; Wang et al., 2021). One major advantage of FL that does not require clients to upload their data which may contain sensitive personal information. Instead, clients separately train local models on their private datasets, and the resulting locally trained model parameters are infrequently synchronized with the help of coordinating server (McMa et al., 2017). While the FL framework helps alleviate data-privacy concerns for distributed training, most of existing FL algorithms critically assume that the clients have enough compute and storage resources to perform local updates the entire machine learning model. However, this assumption does not necessarily hold in many modern applications. For example, classification problems with an extremely large number of classes (often in millions and billions) commonly arise in the context of recommender systems (Covington et al., 2016), information retrieval (Agrawal et al., 2013), and language modeling (Levy & Goldberg, 2014). Here, the classification layer of neural network itself is large enough that a typical FL client, e.g., a mobile or 1oTT device, cannot even store and locally update this single layer, let alone the entire neural network.
Split learning (SL) is a recently proposed technique (Vepakomma et al., 2018: Thapa et al., 2020) that naturally addresses the above issue of FL. It splits the underlying model between the clients and server such that the first few layers are shared across the clients and the server, while the remaining layers are only stored at the server. The reduction of resource requirement the clients is particularly pronounced when the last few dense layers constitute a large portion of the entire model. For instance, in convolutional neural network (Krizhevsky, 2014), the last two fully connected layers take 95% parameters of the entire model. In this case, if we allocate the last two layers to the server, then the client-side memory usage can be reduced by 20x. Nonetheless, one major limitation of SL is that the underlying model partitioning leads to an increased communication cost for the resulting framework. Specifically, to train the split neural network, the activations and gradients at the layer where the model is split (refemee to as cut layer) need to be communicated between the server and

---

Under review as a conference paper at ICLR 2022
Synchronization w/ help the server
Codebook codewwords lobels)
Server-side model
Client-side
mini-hutch of cliens datia
Quantizer
Decoder
Loss
Grad.
Conrector
Clients
Server
Gradient w.r.t.
quantized activation
Quantization Layer
Figure Overview of the proposed algorithm. order to reduce the additional communication between the clients and the server, we propose to cluster similar client activations and send the cluster centroids to the server instead. This is equivalent to adding a vector quantization layer in split neural network training. The success of our method relies novel variant of product quantizer and gradient correction technique. We refer the reader to Section 4 for further details.
clients each iteration. The additional message size is in proportion to the mini-batch size well as the activation 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).
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 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 the server. Interestingly, this is equivalent to adding a vector quantization layer in the middle of the split neural network (see Figure for overview of proposed method). Our main contributions
In this paper, we aim make the split neural network training communication-efficee
follows.
We propose end-to-end communication-efioee neural network splitting- aplitting-based for desource-constraing clients (cf. Section 4). The approach employs a novel compression scheme that leverages product quantization effectively compress the activations communicated between the clients and server.
based FL approach
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. order mitigate this problem, we propose gradient correction scheme for backward pass, which plays 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 (ef. Sec- tion 5). Remarkably, we show that our approach allows for up to 490x communication reduction without significant accuracy loss.
We present a convergence analysis for the proposed method (cf. Section 4.3), which reveals trade-off between communication reduction and convergence speed. The analysis further helps explain why the proposed gradient correction technique is beneficial.
Here we note that 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 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 naturally split across two institutions and the data labels are generated on the server. 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, such as SPLITFED (Thapa et al.. 2020). Thus, our method can also leverage existing privacy preserving mechanisms such as differential privacy (Wei et al., 2020) or instance instance-hiding schemes (Huang et al., 2020) to provide formal privacy guarantees.
2

---

Under review as a conference paper at ICLR 2022
2 BACKGROUND AND RELATED WORKS
Federated Learning (FL). Suppose we have a collection of M clients I =
M }.
Each client i E I has a local dataset Dj and corresponding empirical loss function (:)w) ΣçED. () where w denotes the model parameters, and n, = [Dy denotes the number of samples of the local dataset. The goal of FL is to find a shared model w that can minimize the averaged loss over all clients, defined as follows:
M
F(w) = Σ FiFi(w)
(I)
where p; = =nn / n, is the relative weight of local loss F, Motivated by the data privacy concerns, under a FL framework, the clients perform local training and only communicate the resulting models to 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. Compared to these vanilla FL algorithms, split learning-basee based approaches enables learning on resource-constrained clients at a cost of additionally exposing the label information of clients the server.
M
Large Model Training in FL. To deploy large machine learning models on resource-constrained clients, 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) each client to compute the the client update. This may not be feasible for many clients as the parameters in the classification layer a model can dominate the total number of model parameters (Krizhevsky, 2014).
Split Learning (SL). Split learning (SL) is another way of minimizing (1) without explicitly sharing local data clients (Vepakomma 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 w.) are shared across clients and the server, while the remaining layers (called server-side model parameterized by w;) are only stored and trained on the server. Under this splitting, the original loss function for data sample can be re-written as follows:
[w: E) hunn, uww:: E)), AA E I
(2)
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 baseline in Section 3.
Communication-effcieett SL. and Model Parallel Training. Recently, He et al. (2020); Han et al. (2021) proposed to add 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 store the entire classification layer.
Another relevant approach is reducing communication cost in model parallel (MP) training (Dean et al., 2012). In MP training the worker nodes (dedicated machines instead of resource-constrannes client devices) also need to transfer activations of intermediate layers with each other. Gupta et al. (2020) proposed to sparsify the activations according to the magnitude of each element and achieved up to 20x compression ratio (uncompressed/coomressee)) However, in practice it might be tricky to set the threshold for sparsification. Our proposed method explores an orthogonal dimension of sparsification by clustering similar activations within the mini-batches and communicating only the cluster centroids. Empirical results show that the proposed method can achieve up to 490x compression ratio without significant drop in the model performance.
3

---

Under review as a conference paper at ICLR 2022
Product Quantization. Product quantization (PQ) is 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 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 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 also effective in compressing the output of intermediate layers and present a simple technique to enable backpropagation through the corresponding quantization layer.
BASELINE: THE SPLITFED ALGORITHM
In this section, we review representative FL algorithm based on split learning. namely SPLITFED (Thapa et al., 2020), that provides a baseline training approach in resource-constraingg setting. Each iteration of SPLITFED contains four steps detailed as follows. Wherever it simplifies the presentation, we assume that the mini-baich size B on each client is 1 as the derivations hold true for arbitrary batch sizes.
Client Forward Pass: Let S be a randomly selected subset of clients. For í E S, compute the output of the client-side model z; u(w.:: E), where ID; is randomly chosen training sample, and 2; E Rd denotes the activations of the last layer of the client-side model. Then, each selected client sends 2, (together with its corresponding label if necessary) the server.
Server Update: The server treats all the activations (also referred to as smashed data) (zi)ies inputs perform one step of gradient descent the server-side model: Aw, "h Eies Pt V h (w., / Eies Pt, where 7/s denotes the server learning rate.' In addition, the server also computes the gradient with respect to the activations Vz, h(w,; 2;) and sends it back to the corresponding client i.
3. Client Backward Pass: Each selected client computes the gradient with respect to the client-side model using the chain rule: g; v = =z, h(,,; z:) Vn, u( u(we;:) E)-
Client-side Model Synchronization: At last, the client-side model is updated by synchronizing
client gradients Aw, 77 Eies Pigs/ Eies Pi. where T) denotes the learning rate for the client-side model.
One can easily validate that, with the above four steps, SPLITFED is equivalent to mini-batel stochastic gradient descent (SGD) with total batch size of B x (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 |wa| + Bd. Since the up-link communication bandwidth is quite limited for clients in federated learning, when the mini- 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). Next, we present a novel vector quantization-bast FL training framework that addresses this limitation of SPLITFED enable communication--fticiaae training in resources-constrainge settings.
4 PROPOSED METHOD: FEDLITE
We now introduce our proposed method, FEDerated spLIT learning with vEcttrr quantization (FEDLITE), which can drastically reduce the up-link communication cost on clients in the SPLITFED algorithm. 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
TThe server learning rate here has different definition from previous FL. literature, e.g., Reddi et al. (2020).
4

---

Under review as a conference paper at ICLR 2022
(i) Divide each vector into ( subvectors
(ii) Stack subvectors into R groups
(iii) Assign each subvector to a cluster centroid
Original activations
Quantized activations Z
1.00
1.0.0.100
K-means
g R
Codebook
Codewords
U.0.1.1,0.
K- K-means
0.0.1.0,11
B
gB/R
Communicated
B
Figure 2: Illustration of our proposed quantizer. Given mini-batch of activations Z RR there are three steps: (i) divide each activation vector into q subvectors; (ii) stack subvectors into R groups based on their indices; (iii) perform K-means clustering to get I, centroids for each group. Each subvector can be represented by the index the closest centroid in its corresponding group. On the server, by simply rearranging centroids, we can get the quantized activations Z.
mini-batch are nearly identical to each other, then instead sending 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 the above observation, we let each client input their activations into clustering algorithm to get L centroids. Then, each activation vector is represented by the index its closest centroid. Instead of sending mini-batch of B vectors, only L centroids and the cluster assignments transmitted to the server. This procedure equivalent to vector quantization. Here, a simple choice of the clustering algorithm K-means (Krishna & Murty, 1999). However, as we observe in Figure vanilla K-means clustering leads to a high quantization error with very limited compression ratio. non-trivial to design proper clustering algorithm that can minimize the communication between clients and server, while maintaining a low quantization error. In Section 4.1, we design novel variant of product quantizer that can flexibly operate this trade-off.
While the above approach is appealing, 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 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 critical for our method achieve 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 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.
4.1 FORWARD PASS: COMPRESSING ACTIVATIONS WITH PRODUCT QUANTIZATION
In this subsection, we first present our proposed qunatizer and then elaborate on its advantages. A visual illustration provided in Figure
The Proposed Compression Scheme. Suppose a client has computed batch of activations Z [x(1) 2 2(B) E R using the client-side model we, where d is the dimension of
each activation and B denotes the mini- batch size on each client. Our proposed quantizer first divides each activation vector 20), E [1, B] into q subvectors with equal dimension d/g. We

---

Under review as a conference paper at ICLR 2022
Quantization erro
Durs (1-4k-c)) (3<Reeg)
0.10
nerozse
K-meens (Amqu=1)
0.08
0.06
0.04
0.02
Decreaso
0.00
10a
10'
10'
Compression ratio
Figure 3: Quantization error (lower is hetter) versus compression ratio (larger is better). Our proposed quantizer achieves a better quantization error vs. compression ratio trade-off as compared vanilla product quantization (PQ) and K-means. In this example, the activation size is d = 9216, and the mini-batch size B 20. The activations are trained via two-layer CNN on Federated EMNIST. In each curve, we vary the number of clusters L. For green curves, the number of subvectors takes value in {288, 1152, 4608); for red curves, the number of groups R takes value in (2304, 1152, 384, 1}. and y is fixed as 4608.
denote the s-th subvector of 2(5 as ((.+) E Rd/0, Then, the quantizer stacks all subvectors into R groups based on their indices. For example, the first group g()) contains all subvectors with indices (j, (j, (j.2), 2). (/ R), E [1, Bl: the second group g(2) contains all subvectors with indices (j.g/Rt R + 1), (..// + 2)..... (j. 2g/ R), Vj E [1, B), and forth. At last, K-means clus- tering is performed
subvectors within each group g(r), Vr [L, It] to find L cluster centroids
((r) (c(r.1) c(r,)) c Rd/9, Now each subvector 26,)) can be approximated by its closest
centroid in the corresponding group. Formally, if 2(i+)) g(r), then
g(ia) A clrdo where ((2-0) arg min || c
(E(1,L]
In the above equation, g003) represents the quantized version of 2 ((++), By concatenating all quantized subvectors, server can get the quantized activations Z and use as input to update the server-side model. In order obtain the quantized subvectors, each client needs send all cluster centroids {c("))rE(I R) (referred to as the codebook) and the cluster assignments (?(:.>) [jE(I,B],.E[[ (referred to as the codewords) to the server. Assuming that each floating point number occupies o bits?, the transmitted message size per client is reduced odRL/4 + Bq log, L from ødB.
The Benefit of Subvector Division: Reduced Quantization Error. The first key component of our proposed quantizer subvector division, i.e., setting the number of subvectors q to be greater than 1. When q 1, the quantizer reduces to vanilla K-means. In this case, the codebook size is ødL and each activation vector have L possible choices. When we set q = R > 1, the codebook size still LdL. However, each activation vector now has L possible choices, as each of subvectors has L choices. Thus, the number of quantization levels becomes exponentially larger without any increase in memory usage and computation complexity. As illustrated by the green lines in Figure 3, having more quantization levels or effective centroid choices can significantly lower the quantization error.
The Benefit of Subvector Grouping: Improved Compression Ratio. By using subvector division, although the quantization error can be reduced, the codebook size ødL is still pretty large. This is because subvectors within one vector are quantized separately using different codebooks. To further reduce the communication size, we propose subvector grouping and force subvectors within one group to share the same codebook. As result, the total codebook size can be reduced odLRR/g /q. When > R 2 1, there can be an order of magnitude increase in the compression ratio, as illustrated by the red lines in Figure One can also observe that, by changing the value of R, our proposed trade-off than K-means and
quantizer (red lines) achieves a much better crro-vveruu-compprsss vanilla production quantization scheme.
Why not Reuse the Codebooks from Previous Iterations? In our proposed scheme, the codebook reconstructed at each iteration on clients. This necessary, since in FL, only few stateless clients are selected to participate training in at each iteration and they have non-IID data distributions (Kairouz et al., 2019). Previous codebooks can be stale and other clients' codebooks may not be suitable.
When computing the compression ratio in this paper, always assume Φ 64.

---

Under review as a conference paper at ICLR 2022
Now that we have introduced our product quantization- based compression scheme, next address the issue of performing backpropagation in the presence of the additional quantization layer.
4.2 BACKWARD PASS: GRADIENT CORRECTION
As mentioned at the beginning of Section 4, there is a gradient mismatch problem in the client backward pass. While client needs Vz, h( w' zi) to update the client-side models, it can only receive Vg, h x,) from the server. A naive solution is to treat Vz,h as an approximation of V h and update the client-side model by V h(w,, V V m. u(un;; E) = (a. (00//22))a (0z, However, due to gradient mismatch, this approach can lead to significant performance drop, as observed in our experiments (cf. Section 5). In order to address this issue,
propose gradient correction technique. In particular, we approximate the gradient Vz, h W z;) by its first-order Taylor series expansion: Vz, h( w.; _,) + V h( w.,; _ (z; E.). While the higher-order derivative may be expensive to compute, we use scalar parameter AI > 0 to replace V3 h (w; 3),. Consequently, the gradient of the client-side model is defined as follows:
ah(w,: x
(u(wi:)))
g:
A(z: 2,)
(4)
0z:
Ow.
In Section 4.3, we will provide convergence analysis that can help explain the effects of gradient correction; in Section 5, we will empirically show that setting a strictly positive is crucial for the success of the proposed method. In the following discussion, provide an intuitive explanation of the correction to further motivate our approach.
The Effect of Regularization. We show that the client-side gradient in (4) can be considered gradient of the following surrogate loss (proof is provided in the Appendix):
1
2 2 2
(5)
where u(wei {) is the activation of the client-side model. Besides, 2, 2 -v,h-u, //2 and the quantized activation & are fixed vectors that do not have derivatives when computing the gradient. Setting À > O is equivalent to having regularization term, as per (5). The regularizer encourages the client-side model w. to move in a direction that can decrease the quantization error ||z; ž: Interestingly, using larger value of À may let clients put more emphasis on minimizing the quantization error and lead to quantization--rient
activations. However, one cannot set a arbitrarily large value for X because the client-side model may output the same activation for all inputs and fail to minimize the original loss function.
4.3 CONVERGENCE ANALYSIS
In this subsection, we provide a convergence analysis for FEDLITE. The analysis will highlight how the quantization error influences the convergence and how the gradient correction technique helps. In particular, the analysis is conducted under standard assumptions for mini-batch SGD. We assume the objective function F(w) F(ww,: ux,)) is T--Lipschitz smooth (i.e., F(w) F( < L. 1|w vl)), and stochastic gradient (ww) has bounded variance [ V F( < / where B is the mini-batch size per client and S is the number of selected clients per iteration. Under these assumptions, we have the following theorem.
Theorem (Convergence of FEDLITE). If the client-side model and server-side model update using the same learning rate a = = BS/T where T is the number of total iterations, then the expected gradient norm of the function ming (VF(((()))) II? can be bounded by
4(F(w(()) Fin)
4Lo²
BS T +2
A
+ (A2 (A2
(6)
VBST
VBST
Opt. mini batch SGD
Addul. caused quantization
where k. is the maximal quantization error max || -1I. X is the tunable parameter in our gradient correction scheme, Fing the lower bound of function value, and constants A1, A2, Ag are the largest eigenvalues of matrices 82 hiw,, z)/dziw, 82t Ph =)/ Du(w,; E)/ur, ww, respectively.
7

---

Under review as a conference paper at ICLR 2022
Theorem guarantees that the proposed algorithm FEDLITE converges to neighbourhood around the stationary point of the global function. And the size of this neighborhood in proportion to the maximal quantization eπror k during the training. When there is quantization, the error bound (6) recovers that of mini-batch SGD. Moreover, one can observe that setting a positive À can help to reduce the additional error caused by adding the quantization layer.
5 EXPERIMENTS
We implement the proposed method FEDLITE using FedJAX (Ro et al., 2021) and evaluate its effectiveness three standard federated datasets provided by the Tensorflow Federated (TFF) package (TFF, 2021): ): (i) image classification on FEMNIST, (ii) tag prediction on StackOverflow (referred as SO Tag) and (iii) next word prediction on StackOverflov (referred as SO NWP). On FEMNIST, the same model architecture as (Reddi et al., 2020) is adopted. We place two convolutional layers and two dense layers on the clients and the server, respectively. With this splitting, the client- side model only has about 1.6% trainable parameters of the entire model. Therefore, the client-side resource requirement is significantly reduced. On SO Tag, both the client- and server-side models contains only one dense layer. On SO NWP, we place one LSTM layer and one dense layer on the clients, and another dense layer on the server. The ratios between the client-side model size and entire model size are 83% and 79% on SO Tag and SO NWP, respectively. Here, we note that even though the client-side model sizes for SO Tag and SO NWP do not correspond to an ideal setting for the split learning- baased approaches, we include these two datasets showcase the utility of our proposed method for language-basee tasks.
For each task, we select the learning rate that is best for the baseline SPLITFED algorithm. Although separately tuning the learning rate for our proposed method can further improve its performance, using the same learning rate SPLITFED already demonstrates the advantages of our method. Unless otherwise stated, the number of groups R in our proposed quantizer is set to 1, as it exhibits the best trade-off in Figure 3 and experiments.
Main Results: Effectiveness of FedLite. 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 g 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 10x compression ratio with almost no accuracy loss. Furthermore, we allow for 5% relative loss compared to SPLITFED, then our method can achieve 490 compression ratio 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 5 even improves when adding the quantization layer. We conjecture that compressing the outputs of a intermediate layer might have similar effects to dropout.
Overall Communication and Computation Efficiencies. Here, we provide concrete example showing how the proposed method improves the communication and computation efficiencies over previous works. On FEMNIST, setting g 1152 and L = 2 amounts to 490x compression ratio, which amounts to a 490x 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 10x smaller. Compared to FEDAVG, the up-link communication cost per round reduced by 62x, with 64x 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).
Effectiveness of the Gradient Correction. From Figure 4, it is easy observe that the gradient correction technique (i.c., X > 0) is crucial for improving performance. Without correction, the algorithm can even diverge in the high compression ratio regime. While À is separately tuned for
8

---

Under review as a conference paper at ICLR 2022
Accuracy
Recallies
50
Accuracy
70
15
65
Ours
Ours
Ours
Ours w/o grad. car.
30
Ours weo grad.
10
Ours we grad. cor.
SplitFed
SplitFed
SpiitFad
<5% acc.
<5% drop
5
<5% acc. drop
10a
16a
10a
10
10a
10-
10a
107
10a
10
Compression ratio
Compression ratio
Compression ratio
(a) FEMNIST
(b) Tag
(c) SO NWP
Figure 4: Trade-off between the accuracy and compression ratio. Our proposed method can achieve up to x90x, 247x, and 51x communication reduction with less than 5% accuracy drop on FEMNSIT, SO Tag, and SO NWP, respectively. Each curve in the figures corresponds to one value of g (number of subvectors); and each point on a curve corresponds to a specific value of L (number of clusters). For our method, the number of groups R is fixed as 1.
Accuracy
Accuracy
4608
46000
Subvectors
2304
83.8
83.8
83.8
83.5
82.7
Subvectors
2304
83.7
83.5
83.5
82.8
81.4
83.6
85.5
83:2
03.0
82.0
83.4
83.1
82.6
80.6
80.5
1152
1152
Accura
82.8
02.2
80.9
82.6
80.5
79.5
576
81.0
78.7
76
576
Ours
PQ
288
81.4
79.7
79.8
79.3
77.4
68.3
wio quantszation
144
80.
79.7
78.5
72.8
42.1
144
76.7
74.2
63.0
5.1
5.1
76
<5% acc. drop
L= Clusters
16
2
18
Clusters
10a
Compression ratio
10-
(a) À A= 5 x 10
(b) À =0 0
(c) Quantizer comparison
Figure 5: Ablation studies on FEMNIST. (a) and (b): Validation accuracy on FEMNIST when fixing and varying g and L. Setting a small positive value for À improves accuracy for almost all (g. L) pairs; (c) Thanks to subvector grouping, our proposed quantizer can achieve an order of magnitude larger compression ratio as compared to vanilla product quantization scheme.
each point (i.e., each (q, L) pair) in Figure 4, we found that fixing À for all (q, L) pairs still leads to significant improvements. In Figure 5, report the performance FEMNIST for various choices of and L. In particular, with 288, the accuracy improvement ranges from 3 - 72%. Furthermore, we found that small ranging from 10-3 10-3 works well all training tasks in our experiments. Setting larger À leads to near-zero quantization error. But the model may have poor performance, as tends to output the same activation for all inputs and becomes meaningless.
Effectiveness of the Proposed Quantizer. In Section 4.1, we already show that stacking subvectors into groups (i.e., setting < q) is critical reach a high compression ratio. But how does affect the model performance? To answer the question, we run the proposed method with I = > 1 (vanilla PQ), and report the results in Figure 5c. Observe that the proposed method significantly improves the compression ratio with minimal loss of accuracy.
6 CONCLUSION
In this paper, we studied the problem of training large machine learning models in federated learning setting with resource-constrained clients. Due to limitations on storage and/or compute capacity, the clients cannot locally optimize the entire model. This prohibits the usage of previous federated learning algorithms. Split neural network is promising approach to address this issue, as it only assigns the first few (small) layers of a neural network to the clients. However, the network splitting incurs additional communication. In order to make split neural network training to be communication- efficient, we propose an end-to-end training framework FEDLITE that can compress the additional communication by up to x90x with minimal loss of accuracy. The success of our method relies on variant of product quantization scheme and gradient correction technique. We perform theoretical analysis as well as extensive experiments on both vision and language tasks to validate its effectiveness.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Saurabh Agarwal, Hongyi Wang, Shivaram Venkataraman, and Dimitris Papailiopoulos. On the utility of gradient compression distributed training systems. arXiv preprint IXXv:2103.005433 2021.
Rahul Agrawal, Archit Gupta, Yashoteja Prabhu, and Manik Varma. Multi-label learning with millions of labels: Recommending advertiser bid phrases for web pages. In Proceedings of the 22nd international conference on World Wide Web, pp. 13-24, 2013.
Ting Chen, Lala Li, and Yizhou Sun. Differentiable product quantization for end-to-end embedding compression. In International Conference on Machine Learning, pp. 1617-1626. PMLR, 2020.
Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations. In Proceedings of the 10th ACM conference on recommender systems, pp. 191-198, 2016.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Andrew Senior, Paul Tucker, Ke Yang, Quoc V Le. et al. Large scale distributed deep networks. In Advances in neural information processing systems, pp. 1223-1231, 2012.
Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient feder- ated learning for heterogeneous clients. In International Conference Learning Reprexentations, 2020.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Tiezheng Ge, Kaiming He, Qifa Ke, and Jian Sun. Optimized product quantization for approximate nearest neighbor search. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 2953, 2013a.
Tiezheng Ge, Kaiming He, Qifa Ke, and Jian Sun. Optimized product quantization. IEEE transactions on pattern analysis and machine intelligence, 36(4):744-755, 2013b.
Vipul Gupta, Dhruv Choudhary, Ping Tak Peter Tang, Xiaohan Wei, Xing Wang, Yuzhen Huang, Arun Kejariwal, Kannan Ramchandran, and Michael W Mahoney. Fast distributed training of deep neural networks: Dynamic communication thresholding for model and data parallelism. arXiv preprint ar wrXiv:2010.08890) 2020.
Dong-Jur Jun Han, Hasnain Irshad Bhatti, Jungmoon Lee, and Jaekyun Moon. Accelerating federated learning with split learning on locally generated losses. JCML'21 workshop on federated learning for User Privacy and Data Confidentiality, 2021.
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cnns at the edge. Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020.
Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas D Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. arXiv preprint arXiv:2102. 13451, 2021.
Yangsibo Huang, Zhao Song, Kai Li, and Sanjeev Arora. Instahide: Instance-hiding schemes for private distributed learning. In International Conference on Machine Learning, pp. 4507-4518 PMLR, 2020.
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2010.
Peter Kairouz, Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv:1912.04977,.222 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 14/2.6980, 2014.
10

---

Under review as a conference paper at ICLR 2022
K Krishna and M Narasimha Murty. Genetic k-means algorithm. IEEE Transactions on Systems, Man, and Cybernetics, Part B (Cybernetics), 29(3):433 4399, 1999.
Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv preprint 2014.
Omer Levy and Yoav Goldberg. Neural word embedding as implicit matrix factorization. Advances in neural information processing systems, 27:2177-2185, 2014.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. CoRR, 2019.
Brendan McMahan, Eider Moore, Daniel Ramage. Seth Hampson, et al. Communication-efficennt learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, 2017.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295, 2020.
Amirhossein Reisizadeh, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaie, and Ramtin Pedarsani. Fedpaq: A communication- effircient federated learning method with periodic averaging and quanti- zation. arXiv preprint 0rXiv:1909.130/1 2019.
Jae Hun Ro, Ananda Theertha Suresh, and Ke Wu. Fedjax: Federated learning simulation with jax. arXiv preprint arXiv:2108.00 17, 2021.
Daniele Romanini, Adam James Hall, Pavlos Papadopoulos, Tom Titcombe, Abbas Ismail, Tudor Cebere, Robert Sandmann, Robin Roehm, and Michael A Hoeh. Pyvertical: A vertical federated learning framework for multi-headed splitnn. arXiv preprint 1XXiv:2004000489 2021.
TFF. TensorFlow Federated: machine learning on decentralized data.
https: //www.
tensorflow.ore
federated, fed.r//a.rrrreedrrddddd 2021. Accessed: 30 September 2021.
Chandra Thapa, Mahawaga Arachchige Pathum Chamikara, and Seyit Camtepe. Splitfed: When federated learning meets split learning. arXiv preprint arXiv:2004. 12088, 2020.
Praneeth Vepakomma, Otkrist Gupta, Tristan Swedish, and Ramesh Raskar. Split learning for health: Distributed deep learning without sharing raw patient data. arXiv preprint 1rxiv::8// 18/2.0054, 2018.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan AI-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide to federated optimization. arXiv preprini D.Xiv:2107.0699 2021.
Kang Wei, Jun Li, Ming Ding, Chuan Ma, Howard H Yang, Farhad Farokhi, Shi Jin, Tony QS Quek, and Vincent Poor. Federated learning with differential privacy: Algorithms and performance analysis. IEEE Transactions on Information Forensics and Security, 15:3454-3465 2020.

---

Under review as a conference paper at ICLR 2022
A DETAILED EXPLANATION FOR THE EFFECTS OF REGULARIZATION
In Section 4.2, we argue that the client-side gradient is gradient of the following surrogate loss function
s(w.) - 2 2; 1z
1 -
where 2; - Vz, h(w,; /2 and the quantized activation ス, fixed vectors that do not have derivatives when computing the gradient. Here, we are going to provide a formal proof to support this argument. Directly taking the derivative, we have
@s(w.) Om,
0zi
dzi
-2 21) Sw.
+ A(zi - zzi)
(8)
Sw.
Ozi
(V h 2,) A(z; 2,)) ow. The above equation exactly the same as the client-side gradient we defined in (4).
(9)
B PROOF OF THEOREM
B. PRELIMINARIES
Without loss of generality, we define w (wee wr). Then, the gradient of the original loss function can be decomposed into two parts:
F (w) V m. F(w); F..FF(w))
(10)
Similarly, we denote the stochastic gradients as follows:
g(w) = (w): (w))
(11)
where [w)] = Vw, F(w) and Elg. (w)) = w., F(w). In addition, we assume the stochastic
gradients with mini-batch size of BS have bounded variance:
E [lg(()) - [F(w)]l BS
(12)
That is equivalent to
o²
E (g)
w. F(w)11 + (gg((w) w F(w))l S
(13)
BS
Then, we denote the client- and server-side gradients in the presence of the quantization layer as follows:
[.(w) (w) + (,w))
(14)
(w) = g. + (w) The concrete expressions of 6, and 8, will be derived later.
(15)
B.2 MAIN PROOF
With the above notation, according to the Lipschitz smoothness of the loss function, we have
B [[F(w((+1)) - (0) < V Vw, ((w (w (ww (() (..
a² E ()))
||²
aLL E
p.. (1o(t)))
(16)
For the first term on the right-hand side of (16), we have
(1) V.n. (9) + ((8.
< V (z) /1)
V Vw. F(w
+ 2
13 18.. (w
(17)
< (1 1 4 m. (1))
2
+ S
5 (20
(18)
2e
2
V (20(00)
||
+ 2
1
(w (1)
(19)
12

---

Under review as a conference paper at ICLR 2022
where (17) uses Young's Inequality, and 0 is arbitrary positive number. We set right after (17). For the third term on the right-hand side of (16). we have
E ()) + )
S2E ge
(")
||
+ 2E 6e (2(((()))
(20)
For the server-side gradients Is. repeat the same process. that
Combining them together, it follows
[ - a 2
/| + t ())
| |(
+ 122aL(x + x
2 g + xxx)
(21)
where
X max E |8.(w))1
(22)
2 max E [6,(1))
(23)
Using the assumption that the stochastic gradient has bounded variance, obtain that
B[ ((w(()) S - --a al 1 - ab) / Lo
BS
+2a2 xx2 + 2 (e + X2).
(24)
When aL S 1/4, we have
+ a²
E F( (++2)) - (1 5 a
VF(w
+ 2a L(x + xx))
BS
"2 (xe + x2)
a
(25)
With minor rearranging and taking the total expectation, it follows that
(0)) ||
4E - ((w
4aLo²
V F
BS + (2 + 4aL) +x3).
(26)
Taking the average from t 0 T obtain that
T-1
((F(w(()) F(w)))
4aLo²
+ (2 4aL) +x²)
(27)
aT
BS
0=
If v BS/T.
then
T-1
1)) ()) FF(w"))
4Lo²
4VBSL VT
+
+2
(x. (x.+x x3). (28)
/BST
VBST
Next we are going to show how Xe, Xs relate to the quantization error. Suppose that ((O) denotes the randomly sampled mini-batch on client Then according to the definition of server-side gradient, we have
Sh(w,; Q(u(we: E)))
ge (w)
(29)
S(0)
[
Ow,
(ES
SEB.
Oh ah(w, Q(u(wa:()) Ow,
ah (w.; ul (we:
(g. +
Σ
[S()
|BB())
Ow,
(Est))
SE
(30)
13

---

Under review as a conference paper at ICLR 2022
where Q represents the quantizer and maps the original activation to its quantized version. As consequence,
ah(w, Q(u(wr::) E)))
(h(w,; (w,; uwwc: E))
(,())
(31)
|
18,
Ow,
Ow.
iEste
SEB!
and
dh( Wsi Q(u(w,:E)) ah(w, u(w::
[[6,(w)]'
max {ED,JEZ
(32)
Ows
Ow,
For the ease of writing, we denote we have
u(we {) and E = Quu(w,) ()). Using mean value theorem,
Sh(ws, u)
[[)((w)] S max
(z (z-2) 2)
(33)
Sudw,
SApnk
(34)
where to + t) for some t 1, constant A₁ is the largest eigenvalue of matrix 02 (2w u//duw,, and x = max 2 =1 denotes the maximal quantization error.
Using the same technique, for the client-side gradients, we have
Oh(un: ((u(wr:()) Oh (ug u uwwe E)) (Q(u(we:E))
Ōu(we: E)
6c(w))
Σ
|800
(B(0))
Ou(wo: E))
Sw.
(ESt))
(EB(1)
@u(weik) Owc
Σ X(u(we):) Q(u((nc:
(35)
[St0]
(B((1) SEB!")
VES())
Accordingly,
Öh(w,;
@h(wsiz)
Oz
[1,,(w))
max
+A(z 2)
(36)
0z
Oz
Omc
[0Ph(ws;ur u)
2) (3 (3- z)
Oz
maax
(37)
Ou
Dur.
<Aa
A3
(38)
where A2 and Ag are the largest eigenvalues of matrices h(w, u)/Ou² and 02/Ow, respectively. Substituting the above bounds (34) and (38) on 82, 5, into (28), we have
T-1 Σ
((((w()) (0) F(w"))
4Lo²
VBST
VBST
AVBS BS
+
+2 (A? + (A2 - 2)2 max ||> - 2112. (39)
VT
Here complete the proof.
C MORE EXPERIMENTAL DETAILS
C.I ADDITIONAL RESULTS
We additionally report the training curves of FedAvg, FedLite and SplitFed on EMNIST in Figure It can be observed that in terms of communication cost, FedLite is significantly faster than other two baselines. We also need to highlight that FedLite and SplitFed have lower memory and computation requirements clients than FedAvg.
14

---

Under review as a conference paper at ICLR 2022
Accuracy
FedAvg
SplitFed
10a
10a
10a
10*
Communication size (MB)
Figure 6: Training curves for different algorithms on FEMNIST. Although split Haaming-hased approaches (SplitFed and ours FedLite) communicate every iteration, they communicate much less than FedAvg
C.2 HYPER- PAREERRR CHOICES
Here, summarize all the hyper-parameteed three training tasks.
FEMNIST
Best learning rate: 10-1.5 115
Optimizer: SGD
Mini-batch size per client B: 20
Activation size d: 9216
Number of clients per iteration: 10
Ranges of g: {4608, 2304, 1152, 576, 288, 144]
Ranges of [32, 16, 8, 4,
Ranges of X: {0, 10 110-0.5550000 10 1,-14,,, 10-4)
Client-side model: Same the first five layers of the neural network used in (Reddi et al.,
2020): Conv2d + Conv2d + MaxPool2d + Dropout + Flatten. Model size: 18, 816 x 64 bits.
Server-side model: Same the last three layers the neural network used in (Reddi et al., 2020): Dense Dropout Dense. Model size: 1, 187, 774 x 64 bits.
SO NWP
Best learning rate: 0.01
Optimizer: Adam (Kingma & Ba, 2014)
Mini-batch size per client B: 128. Each sample contains 30 words. So the effective batch size is 3840.
Activation size d: 96
Number of clients per iteration: 50
Ranges of q: {48, 24, 6,3)
Ranges of L: {960,, 480, 240, 120, 60,
Ranges of X: {0, 5 10
10 10
Client-side model: Same as the first three layers of the neural network used in (Reddi et al., 2020): Embedding + LSTM + Dense. Model size: 3, 680, 360 64 bits.
Server-side model: Same as the last layer of the neural network used in (Reddi et al., 2020):
Dense. Model size: 970, 388 x 64 bits.
SO Tag
Best leaming rate: 10 -0.0.
15

---

Under review as a conference paper at ICLR 2022
Optimizer: AdaGrad (Duchi al., 2011)
Mini-batch size per client B: 100.
Activation size d: 2000
Number of clients per iteration: 10
Ranges of q: (1000, 500, 250, 200, 125, 25)
Ranges of L: (100, 60, 40, 20, 10)
Ranges of A: {0, 10- x 10 3, 10 2,5 10
Client-side model: One dense layer. Model size: 5000 x 2000 x 64 bits.
Server-side model: One dense layer. Model size: 2000 x 1000 x 64 bits.
16


Reviews:
[{"review_id": "SqInKBv5Z35", "reviewer": "Reviewer_oo1D", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Pros**:\n- The paper is well written and structured.\n- The paper provides theoretical insight about FedLite convergence rate.\n- Interesting and novel idea aim to reduce the communication cost.\n\n**Cons**:\n- The authors should cite additional previous works that tried to address the problem of client's device heterogeneity [1,2,3].\n- My main concern is that under supervised learning setup FedLite sends the labels to the server to obtain $\\nabla_{z_i}h\\(w_s;z_i\\)$. Sending the labels to the main server compromises privacy (for example in healthcare systems).\n- I think the authors should add more challenging experiment to check the trade off between compression and accuracy. (Celeb a from LEAF dataset for example).\n\n**References**\n[1] Shamsian, Aviv, Aviv Navon, Ethan Fetaya, and Gal Chechik. \"Personalized Federated Learning using Hypernetworks.\" arXiv preprint arXiv:2103.04628 (2021).\n[2] Achituve, Idan, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. \"Personalized Federated Learning with Gaussian Processes.\" arXiv preprint arXiv:2106.15482 (2021).\n[3] Cho, Yae Jee, Jianyu Wang, Tarun Chiruvolu, and Gauri Joshi. \"Personalized Federated Learning for Heterogeneous Clients with Clustered Knowledge Transfer.\" arXiv preprint arXiv:2109.08119 (2021).", "summary_of_the_paper": "The authors propose end to end training framework that relies on a novel vector quantization.\n\nThe authors claim the following contributions:\n- End-to-end communication-efficient neural network splitting-based FL approach for resource-constrained clients.\n- New gradient correction scheme for backward pass.\n- Reducing communication rounds up to 490X without compromising the accuracy.\n- Providing ", "summary_of_the_review": "Questions were raised during the review process. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "D361STFSZRD", "reviewer": "Reviewer_WW1C", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I truly enjoyed reading this paper. I had thought about the problem before and was excited to read a paper that addresses the communication of split activations in the Split Learning framework. The proposed solution to the problem is presented clearly and the experimental evaluations are thorough. \nRegardless, there are a few questions I would like to see answers, as well as some issues I have with the paper, the removing of which would make it a lot stronger.\n\nIntroduction & Motivation:\n- While Split Learning allows to 'outsource' the final layers' heavy computations to a server, this inherently reduces the privacy compared to FedAvg based approaches that communicate weights-only. Especially since labels have to be communicated for the loss-computation. While I understand that there is a need for SplitLearning regardless of this disadvantage, I would encourage the authors to mention this significant downside to SL in either introduction or discussion.\n\nBackground/RW:\n- The authors mention model parallel training in some places in their paper. However they do not focus at all on it, provide no experiments for it and ignore questions that arise when using their approach in the model parallel training setting. More specifically, in model-parallel training the additional latency of the quantisation approach cannot be ignored and would presumably cannibalise on any advantages that the reduced communication brings with it. Either the authors include a proper analysis of Model Parallel Training or they should remove the discussion around it. \n- There is large body of literature that studies quantised training of neural networks. Some works focus on weight-quantisation, others on activation-quantization (and the combination of both) for the purpose of efficient inference on resource-constrained devices. While many different and sometimes complicated approaches exist, the simplest approaches consider stochastic scalar quantization in combination with the Straight-Through estimator. I am not so familiar with the Split Learning literature, but I would expect this to be a trivial baseline for quantising the activations of a single layer's activations in the forward pass (as considered in this paper). In order to be applicable for sped-up inference in these settings the quantization grid needs to be fixed (optionally learned) at the end of training. In the split learning setting, this grid could be constructed arbitrarily for any chunk of the b x D activation matrix, similarly to the authors' discussion in this paper. Naturally the split-learning approach is not constrained to uniform grids (as would be for inference acceleration on e.g. in8 inputs), so you might even consider quantising a transformation of the activations with e.g. the Normal CDF (assuming activations are normally distributed). I would ask the authors to either consider including experiments for these quantization approaches in their paper or elaborate why this would be out of scope for Split Learning, maybe I am missing something...\n\n\nMethod\n- The authors mention that client-weights $w_c$ need to be synchronised every iteration. Is that the default in SplitLearning or would you also consider averaging every E local epochs, as is done in FedAvg traditionally? \n- You mention that you assume 'each float-number' to require 64bits. In the context of FedLITE that is relevant for the code-book entries. Is that a requirement for your method to function well? Since NN activations are usually 32bits (16bits are often enough), I am surprised by that assumption. \n- Computational complexity: Can you elaborate on how expensive it is to compute the quantised activations for a mini-batch of data? Ideally you would provide timing numbers of forward/backward pass and quantization on a given hardware and discuss the impact on resource-constrained clients. \n- You discuss the infeasibility to share codebooks from previous iterations. Did you consider sharing codebooks between data-points of given class, assuming that activations especially higher up in the network share many similarities conditioned on the label? Since you have to transmit the label anyways, decoding would come for free. \n- Your discussion surrounding Eq. (5) is not entirely clear to me: What is the benefit from considering Eq (4) as a gradient of the loss in (5)? The first term involving $\\hat{z}_i$ is not intuitive to me. \n- One of the core insights of Figure 3 seems to be that decreasing R to be =1 seems to have the most favourable scaling properties. Maybe that insight could be reinforced and pointed towards when setting up that hyper-param in the experiment section. \n- I did not analyse the provided convergence proof in detail\n\nExperiments\n- Hyperparameter selection: How did you end up at the hyper parameters chosen for the base SplitFed algorithm? \n- At the last paragraph of 5.1, the authors compare against FedAvg, which assumes a certain number of local epochs E (unless the authors compare against FedSGD). What is the assumption here?\n- Does the optimal lambda change significantly with different hyper-parameters L and q? (Figure 4)\n- Codebook-design: I would be interested to see the consequences of quantising the codebook to less than $\\phi$ bits. Is that something sensible to consider? Additionally, it would be interesting to study the entropy over the selected code-words: Is the method making use of the available code-words? If not, Entropy-coding of the code-words might be a viable approach for further communication reduction. \n\nSmall things:\n- In the introduction you make a claim about AlexNet, however it sounds like you claim a general CNN architecture.\n- Typo above Eq (2): 'while and the remaining layers'\n- I understand the need for an acronym, but the letter-selection seems random... :D On the other hand I have no better suggestion. \n", "summary_of_the_paper": "The authors propose a method to compress the intermediate activations of a neural network that is split across a client and server in the Federated Learning setting. In this 'Split Learning' scenario, transmitting activations and corresponding gradients incur significant communication costs. The authors discuss FedLITE, a method to compress activations during the forward pass and to correct the gradient to account for the quantization during in the forward pass. Their paper is augmented with an interpretation of the gradient correction term as minimising a surrogate loss function, as well as a convergence analysis of their approach and empirical evaluations. ", "summary_of_the_review": "This is a promising paper and discusses a relevant problem. Some detailed questions remain and some changes I believe are necessary. If the authors address my concerns I will improve my rating of the paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "CJdq2GyJ8C", "reviewer": "Reviewer_cPmB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* The approach is well motivated to reduce the transmission cost in split federated learning. \n* The gradient correction scheme is a principled solution to tackle the noisy gradients issue and has solid theoretical implications. They also analyzed the theoretical error bound of their proposed approach which further echos the role played by the gradient correction term.\n* Their approach shows empirical effectiveness compared with the Split learning baselines. The importance of a positive lambada is clearly analyzed by ablation studies.\n \nWeakness / Questions:\n* This paper addresses a specific federated learning scenario, i.e. the split learning setting. However, there are other FL approaches towards efficient parameter transmission or computation cost. I feel that the narrow problem setting that FedLite aims to tackle somehow weakened its contribution. It is unclear to me how FedLite may or may not benefit other FL approaches.\n* The experiment section is not extensive enough, as only one splitFed baseline is compared. I am curious to see how the asymptotic performance of the FedLite is dropped compared with non-split learning baselines, especially FedAvg. ", "summary_of_the_paper": "This paper proposed FedLite, which extends a specific federated learning setting called split-learning to be more transmission efficient by a vector quantization scheme. In split federated learning, the classification layer that has a larger parameter set is saved and learned in the server, while the client learns the relatively lightweight feature layers. In FedLite, the activation outputs of the cut layer are compressed by vector quantization before sending to the coordinate server. To alleviate the issue of noisy gradients caused by compressed activations, a regularization term is introduced into the gradient update step which is derived from the Taylor expansion. Empirical results show the effectiveness of the proposed approach, which largely reduces transmission costs without significant performance loss. ", "summary_of_the_review": "This paper is a solid extension of the Split Federated Learning work, which improves transmission efficiency with an acceptable performance drop. The approaches are well motivated and theoretically verified, however the contribution scope of this work is kind of limited and the empirical study part can be further enriched by discussing and comparing with more related work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Bjj2R91vyyG", "reviewer": "Reviewer_ELQm", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper proposes to use vector quantization to enhance the communication efficiency of SL, a popular paradigm to conduct federated learning. Both theoretical and empirical results are provided to show the effectiveness.\n\n\nPros:\n1. The paper is well written, and the research direction of enhancing the communication efficiency of federated learning is potentially impactful. \n2. The proposed method is analyzed under both theoretical and empirical lenses. \n3. The detailed hyper-parameters and training schedules are provided, which gives good evidence of the reproducibility of the experiments. \n\nCons: \n1. The main concern I have for the current version of the draft is that the proposed FedLite method is not compared against any baseline methods (except for SL itself). Let’s assume we want memory efficient and communication efficient federated learning method, an immediate idea is to use EfficientNet/MobileNet (which are both small and computationally efficient) with FedAvg. How does FedLite look like compared against that vanilla baseline? \n2. The paper proposes to use vector quantization without motivation. It is pointed out that the sparsification-based method requires setting the targeted sparsity, but other than the sparsification method, there are still a series of methods to compress tensors, e.g., [1-5]. The authors are expected to at least motivate the proposed method a bit more. \n3. I wonder if the communication inefficiency of SL is only caused by the CNN and LSTM architectures. In Transformers, splitting Encoder blocks (which is widely used in pipeline parallelism [6-7]) usually leads to quite good communication efficiency. If so, does it mean FedLite can only attain marginal gains on Transformer based NNs? \n4. I’m highly curious about the convergence curve of FedLite compared against FedAvg. It seems FedLite is a FedSGD-ish method, which requires a lot more rounds of communication compared against FedAvg (as stated in the very first FedAvg paper). It is widely known that the communication cost depends on message size and latency (which is related to the round of communication). I wonder if FedLite does lead to significant wall-clock time-saving. \n5. One potential approach to reduce the required number of communication rounds for FedSGD-like methods is to increase the batch size. However, since FedLite aims at low-memory training, it limits the space of using large-batch training. \n6. A recent technical report indicates that the super-high compression ratio is actually meaningless with reasonable bandwidth [8]. In federated learning applications, the bandwidth is usually limited. Thus, it would be useful to show the actual communication time saving of FedLite. \n7. It is stated that FedLite can be combined with existing privacy-preserving methods. But that aspect is not discussed in the paper at all. It’s important for the authors to discuss how to make FedLite private since the vanilla version clearly discloses the users’ privacy. \n\n[1] https://arxiv.org/abs/1802.04434\n\n[2] https://arxiv.org/abs/1610.02132\n\n[3] https://arxiv.org/abs/1905.13727\n\n[4] https://arxiv.org/abs/1806.04090\n\n[5] https://arxiv.org/abs/1705.07878\n\n[6] https://arxiv.org/abs/1806.03377\n\n[7] https://arxiv.org/pdf/2104.04473.pdf\n\n[8] https://arxiv.org/pdf/2103.00543.pdf", "summary_of_the_paper": "This paper introduces FedLite, a compression method built on top of Split Learning (SL) for better communication efficiency. Under the cases where the size of the intermediate output is large, SL can introduce a communication bottleneck in model training. FedLite tackles that issue via compressing the intermediate output via a vector quantization-based method. A gradient correction scheme is proposed to further enhance the proposed compression method. Convergence analysis and experimental results are provided to justify the effectiveness of FedLite.", "summary_of_the_review": "Frankly, I think the paper was rushed out. The idea is sweet but clearly requires more effort to show its effectiveness. The current draft only comes with very preliminary experimental results without comprehensive comparisons against various baseline methods. The motivation of the proposed vector quantization method is also not clear. Overall, I do not think the current draft deserves to appear for a presentation at ICLR 2022.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a compression method for distributed Split Learning for better communication efficiency, by compressing the intermediate output between client and server model by vector quantization. Convergence analysis and experimental results are provided. 
Unfortunately consensus among the reviewers remained that it remains slightly below the bar after the discussion phase. Main remaining concerns were the variety of baselines and benefits from split learning setup in experiments, compared to other FL approaches, quantization approaches, architecture splits. Reviewers also missed a discussion of latency requirements of model-parallel training in FL as opposed to data parallel which allows less frequent communication compared to here (e.g. discussing the split layer size vs latency trade-off, here of quantized intermediate layers compared to regular FL). The newly added Figure 6 does not specify or vary the number of local steps (or batch size) in FedAvg.
We hope the detailed feedback helps to strengthen the paper for a future occasion.
Decision: Reject