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

1 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 is that it 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 a coordinating server ( McMahan 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 on 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 a neural network itself is large enough that a typical FL client , e.g. , a mobile or IoT device , can not 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 at the clients is particularly pronounced when the last few dense layers constitute a large portion of the entire model . For instance , in a 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 20× . 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 ( 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 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 ) . 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 an end-to-end communication-efficient neural network splitting-based FL approach for resource-constrained clients ( cf . Section 4 ) . The approach employs a novel compression 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 for the backward pass , 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 . Sec- tion 5 ) . Remarkably , we show that our approach allows for up to 490× communication reduction without significant accuracy loss . • 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 . Here we note that our proposed method has potential applications beyond FL with resourceconstrained clients as it can be applied in any learning framework that can benefit from a 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 . 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-hiding schemes ( Huang et al. , 2020 ) to provide formal privacy guarantees . 2 BACKGROUND AND RELATED WORKS . Federated Learning ( FL ) . Suppose we have a collection of M clients I = { 1 , 2 , . . . , M } . Each client i ∈ I has a local dataset Di and a corresponding empirical loss function Fi ( w ) =∑ ξ∈Di f ( w ; ξ ) /ni , where w denotes the model parameters , and ni = |Di| 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 : F ( w ) = M∑ i=1 piFi ( w ) ( 1 ) where pi = ni/ ∑M i=1 ni is the relative weight of local lossFi . 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 can not fit in clients memory or require large compute . Compared to these vanilla FL algorithms , split learning-based approaches enables learning on resource-constrained clients at a cost of additionally exposing the label information of clients to the server . Large Model Training in FL . 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 ) . Split Learning ( SL ) . 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 wc ) are shared across clients and the server , while the remaining layers ( called server-side model parameterized by ws ) are only stored and trained on the server . Under this splitting , the original loss function for a data sample ξ can be re-written as follows : f ( w ; ξ ) = h ( ws ; u ( wc ; ξ ) ) , ∀ξ ∈ Di , ∀i ∈ I ( 2 ) where u is the client-side function mapping the input data to the activation space and h is the serverside 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 . Communication-efficient SL and Model Parallel Training . 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 . 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-constrained 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 20× compression ratio ( uncompressed/compressed ) . 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 490× compression ratio without a significant drop in the model performance . Product Quantization . 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 highdimensional 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 .

Summary: 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.