Adaptive Sparse Federated Learning in Large Output Spaces via Hashing
Anonymous Author(s)
AfﬁliationAddress email
Abstract
This paper focuses on the on-device training efﬁciency of federated learning (FL),and demonstrates it is feasible to exploit sparsity in the client to save both computation and memory for deep neural networks with large output space.
To this end,we propose a sparse FL scheme using hash-based adaptive sampling algorithm.
Inthis scheme, the server maintains neurons in hash tables.
Each client looks up asubset of neurons from the hash table in the server and performs training.
With thelocality-sensitive hash functions, this scheme could provide valuable negative classneurons with respect to the client data.
Moreover, the cheap operations in hashingincur low computation overhead in the sampling.
In our empirical evaluation, weshow that our approach can save up to 70% on-device computation and memoryduring FL while maintaining the same accuracy.
Moreover, we demonstrate thatwe could use the savings in the output layer to increase the model capacity andobtain better accuracy with a ﬁxed hardware budget.
Introduction
Recently, federated learning (FL) [18] and its applications [36, 14, 17, 23] receive attentions fromboth research community and industry.
FL deﬁnes a practical yet challenging task: given a set ofdevices where each device maintains its private data locally, we would like to collaboratively traina model on these devices without data exchange.
Signiﬁcant effort has been made in improvingoptimization strategy [15, 31], privacy protection [8] and fairness [16] of FL.
A Challenge in On-device Training: FL introduces a device shift in the distributed training ofmachine learning models.
In the cloud center, we were able to train large-scale foundation models onmassive graphic processing units (GPUs) in a centralized way.
In FL setting, our training hardware islimited to system-on-chip (SoC) on mobile devices.
This shift leads to signiﬁcant efﬁciency issuesin FL: (1) The models we trained on GPU clusters are giant in terms of parameters.
It is standardto have billion-scale parameters for language [2] and recommendation [19] models.
However, thememory constraint for FL devices forces us to limit the model parameter size to make on-devicetraining feasible, which causes a model size mismatch between FL and centralized models.
Thismismatch would degrade the FL model performance and prevent us from the beneﬁts of large deepmodels.
In the centralized training, the advantages of specialized hardware such as TPU provideefﬁcient matrix multiplication for training deep neural networks.
However, the on-device tensor chipslike TPUs in Pixel phones [24] are not as powerful as their serverside counterpart, which wouldsigniﬁcantly improve the training efﬁciency in FL.
The training of deep models through forwardand backward propagation requires memory to store the intermediate results.
Same the previous twoissues, the hardware constraint of mobile devices would affect the efﬁciency of memory read andwrite during training, which exaggerates the computation overhead in FL.
An Opportunity of Sparsity: There is an emerging trend on exploring sparsity in neural networktraining [27, 7, 6, 10].
An interesting direction is to switch the matrix multiplication from denseto sparse mode.
For instance, given an input matrix X ∈ R n × d and the linear layer weight matrix
W ∈ R d × m , we view each row of X as an embedding and each column of W as a neuron.
In this way,for each embedding in X , we could only select a subset of neurons in W for computation.
As a result,we perform a sparse version of operation XW .
Well-known research literature in this area includesthe lottery ticket hypothesis [12, 37], sub-linear deep learning engine (SLIDE) [7, 6] and independentsubnet training [38].
In this paper, we argue that there is an opportunity to improve the on-devicetraining efﬁciency of FL with sparsity.
Firstly, although the sparse training strategy does not changethe model architecture, it only activates a subset of model parameters in each iteration.
This featurecould help us in FL so that each client only selects a subset of trainable parameters from the modelfor iterative optimization.
As a result, the on-device parameter size would be reduced.
Moreover, thesparse alternative to the matrix multiplication could signiﬁcantly reduce the computation overhead,making FL easy on CPU only devices.
Furthermore, the sparse training generates sparse intermediateresults, which also reduces the memory access on the device.
Exploiting Adaptive Sparsity in Large Output Spaces: In this paper, we focus on the sparse FLfor deep neural networks in large output spaces (LOS).
LOS is common in deployed deep models.
Forinstance, in language processing tasks such as next word prediction [20] and question answering [21],the output space would be the vocabulary size.
In recommendation systems [19], the output spacewould be the number of products in the database.
In both cases, the number of classes in theoutput space could be enormous.
As a result, the output linear layer would contain the most modelparameters if we would like to train them on-device using FL.
On the other hand, since we wouldperform Softmax function on the output logits, we could approximate the output layer by focusing onthe logits with high values.
In fact, the LOS would be a perfect scenario for sparse training.
If wecould adaptively pre-select the neurons in the output linear layer that may incur a large inner productwith the hidden input vector, we could only do forward and backward computation on the selectedneurons [1, 7, 10, 6, 30].
Therefore, we could perform efﬁcient on-device FL by saving both thecomputation and memory.
However, the combination of sparse training with FL could still be challenging: (1) An efﬁcientdesign is required for sparse training in FL so that we can maintain the full model on the server anda sparse model on the client device.
(2) It remains unknown whether the adaptive sparsity wouldbe effective in the federated optimization with non-i.i.d data distribution.
(3) How to use the savedcomputation and memory by sparsity for further improvements in the model accuracy.
In other words,how to improve the model performance with a ﬁxed hardware budget using sparsity?
1.1 Our Contributions
In this paper, we introduce an empirical study on the sparse FL in LOS.
We propose to use hashingalgorithms that adaptively select neurons in the output layer for forward and backward computation.
Speciﬁcally, our contributions could be summarized as:
We introduce an adaptive hash-based sparse FL scheme for training in LOS.
In this scheme,the server hashes the output layer’s neurons in hash table.
Next, the server sends the hash
function to each client.
Each client uses the hash function to generate hash codes of their
own data.
Next, each client uses the hash codes to look up the near neighbor neurons of its
data from the server.
Finally, each client only performs forward and backward propagation
on the selected neurons.
We empirically show that the hash-based sparse training in the output layer is effective in the
federated optimization.
We could maintain the same model accuracy with 30% parameters
in the output layer.
As a result, the on-device training efﬁciency would be improved.3.
We demonstrate that in the proposed hash-
based sparse FL scheme, the saved on-devicemodel parameters in the output layer would be used to improve the model capacity.
Weshow that on a ﬁxed on-device parameter budget, if we perform sparse training on the outputlayer and use the saved parameters to increase the embedding and hidden dimension of themodel, we could have better accuracy with on-device FL.
Related Work
Hashing Algorithms for Sparse Machine Learning:
The hashing algorithms have demonstratedempirical effectiveness in the sparse training of machine learning models [7, 10, 6, 34, 32].
[7]proposes SLIDE algorithm that uses locality-sensitive hashing (LSH) [11] to preprocess the neuronsof a wide output layer in hash tables.
Next, given a batch of embeddings, SLIDE use them as a queryand lookup the neurons that are close in cosine similarity from hash tables.
Finally, SLIDE onlyperforms forward and backward computation in the selected neurons.
[10] shows that SLIDE can befurther accelerated with the advance in hardware.
[6] demonstrates that SLIDE could be improved bylearnable hash functions.
[34] focuses on the Frank-Wolfe optimization algorithm.
In particular, [34]preprocess the vertices of the weight space in hash tables.
Next, given the current weight, instead ofcomputing it with all vertices, we only need to compute with near neighbor vertices and choose oneas the next direction.
[32] proposes parallel memory writing algorithms for the sparsiﬁed gradients inthe data-parallel distributed training and provide up to 3 .
52 × speedups.
Efﬁcient On-device Training in FL: There are two major techniques for improving the on-deviceefﬁciency of FL.
The ﬁrst technique is named partial variable training (PVT) [35, 26].
PVT aims atfreezing a fraction of trainable parameters when we train models on the client devices.
For instance,we could freeze some of the fully-connected layers and only perform federated optimization onthe other parts of the model.
[26] focuses on saving the communication cost of transferring modelgradients.
[35] is also trying to save the on-device model size and memory for intermediate results.
Meanwhile, we observe some computation saving in the backpropagation.
Another technique is calledfederated dropout (FedDrop) [3, 9].
The FedDrop randomly selects a subset of neurons from themodel and sends it to the client for training.
Although FedDrop saves the on-device computation andmemory, the nature of randomness would cause an accuracy gap between the FedDrop and originaltraining when we increase the sparsity.
The major reason behind this phenomenon is that FedDrop’srandom sampling is not adaptive to the status of input embeddings.
For instance, it is shown that theneurons with large inner products to input embedding should be a more important example in theoutput layer.
As a result, the missing of these neurons would lead to slower convergence.
Method
In this section, we introduce our hashing algorithms for sparse federated learning in wide output layer.
We start with a formal introduction of hash-based sampling.
Next, we propose a sparse FL schemeusing hashing.
3.1 Hash-based Sampling in Neural Network
In this section, we present how to use hash-based Sampling [27, 33, 7] in the training of neuralnetwork.
We start with introducing a simple yet effective LSH function, namely SimHash [4].
Deﬁnition 1 (SimHash) .
Let K denote the number of hash bits.
Let A ∈ R K ×
d denote a randommatrix where each entry is drawn i.i.d from normal distribution N (0 , 1) .
Given an input vectorx ∈ R d , we deﬁne the SimHash function h : R d → R k ash ( x ) =
sign ( Ax ) , where sign is an element-wise sign function that set nonzero values to 1 and other values to 0.
Moreover, we show that for any two vectors x, y ∈ R d ,
[Equation]
where θ xy is the angle between x and y .
The SimHash’s deﬁnition (see Deﬁnition 1) suggests that if two vectors are close in angle, withhigh probability they would have the same hash code.
Moreover, previous work suggests that witha pair of asymmetric transforms applied on x and y [25], respectively, the collision probability
Pr[ h ( x ) =
h ( y )] would be monotonic to the inner product x
(cid:62) y .
Taking advantages of this property,the hash-based sampling algorithm for a linear layer can be summarized as below:
Build hash hash
Given a batch of embedding X ∈ R n × d , for each row x j in X , compute h ( x j ) and lookupthe W
i s that has the same hash code with h ( x j ) .
Next, we take the union of the retrieved
W i s and write it as matrix W select .
Finally, we subsample on W select with ﬁx sparsity ratioand compute XW select for forward/backward propagation.
In practice, we use L hash tables and take the union of weight columns retrieved by each hash table.
Noted that both K and L are tuning parameters.
For an input batch of embedding X , hash-basedsampling could return a sub-matrix W select that the inner product between each row in X and eachcolumn in W select may incur large inner product.
In LOS, W select represents the neurons that mayhave larger logits with X .
In the practical setting, we may sub-sample on the columns of W select orrandomly add more columns to W select so that we can ﬁx the sparsity budget.
3.2 A Scheme of Sparse FL
In this section, we design a server-client scheme for sparse FL in LOS.
As introduced in Section 3.1,hash-based sampling could select the neurons that may have large logits in the LOS.
However,the maintenance of hash tables and the preprocessing of neurons may generate extra computationoverhead.
In fact, the computation in hashing the neurons should be carefully handled by smartscheduler in centralized training [6].
In our work, we take a FL view of this procedure and arguethat the computation overhead can be reduced by the powerful computation resources in the server.
Speciﬁcally, we introduce our scheme as below:
Hash table maintenance: The server uses SimHash (see Deﬁnition 1) and preprocessesthe columns of the weight W in the output layer in hash table.
The server refreshes the hashtable after each round.
Client initialization: When a new client comes, the server sends the hash function and theweights of layers before the output layer to the client.
The client performs forward pass andgenerate embedding vectors of its own data.
Client hashing:
The client hashes the input embedding using of data samples the receivedhash function and generates a set of bucket locations in the hash table.
Neuron retrieval: The client transfers the bucket locations to the server and looks up theneurons of output layer in the corresponding buckets.
On-deivce training: The client receives the lookup-ed neurons and performs forward andbackward computation.6.
Aggregation:
The client passes the model updates such as gradients back to the server foraggregation.
The advantages of the proposed scheme can be summarized as: (1) the client device does not have tomaintain the hash tables, which reduces the computation overhead in preprocessing neurons duringthe centralized training, (2) the client only lookups a subset of parameters in the output layer, whichreduces the communication cost and on-device model memory, (3) the client trains on neurons thatmay have higher logits in the output layer, which served as an effective negative sampling for fasterconvergence. It is normal that the number of neurons retrieved from hash tables is larger than theclient budget. We can compute the activations of these neurons with the client data and only keep thelarge activation neurons on-device for backpropagation. Note that our approach may involve morecommunication between servers and clients.
We could use the FedSelect [5] to look up neurons withprivacy protection.
3.3 Improving Model Capacity in Fixed Budget
As shown in the previous sections, our sparse FL scheme with hashing reduces the on-deviceparameter size for the last output layer.
For a SoC with ﬁxed hardware memory budget, we coulduse the saved space to increase the trainable parameters in other parts of the model to have betterperformance.
We suggest two major directions: (1) increase the hidden dimension in both embeddingand linear layer for better feature representation, (2) add more attention blocks [29] or linaer layersfor better feature mixing.
In the experiment section, we will discuss how these two directions wouldhelp us improve the empirical performance of on-device FL. [Table]
Experiment
In this section, we introduce an empirical evaluation of the proposed sparse FL training scheme in
LOS with hashing.
We start with introducing the next word prediction task we focus on.
Next, weintroduce the models we evaluate.
Finally, we present the experimental results with an abalationstudy.
We also provide a visualization of hash tables in Appendix A.
4.1 Settings
Dataset.
We evaluate the proposed sparse FL scheme on the next word prediction task using Stack
Overﬂow (SO) dataset.
The SO dataset contains 342477 training clients.
The total training examplesize is 135M. The SO dataset has 38758 clients for validation with dataset size 16M
In the next wordprediction setting, we take the ﬁrst 256 sentences and truncate each sentence to a sequence length 20.
We aim to predict the next word given the previous context words.
It is standard to set the vocabularysize to 10K by taking the most frequent words from the training data.
In our paper, we also extendthe vocabulary size to 80K since larger vocabulary has larger word coverage.
The Transformer Model.
In this paper, we study the performance of Transformer model.
We proﬁlethe parameter size of transformer models in Table 1.
In the model, we set the token embedding size as96.
We also set the hidden dimension of the Q, K, V layer in attention as 96.
For the FFN dimensionin attention block, we set it to 1536.
There would be 330K parameters for each attention block.
Butthe embedding table size would be 960K for 10K vocabulary and 7.68M for 80K vocabulary.
In thispaper, we shared the weights between the embedding and the last output layer so that the largestparameter tensor in the model is the embedding table.
In this case, if we train the full model on eachclient, we have to send all the embedding tables to the client so that they can use them for outputlayer.
On the contrary, if we could perform sparse training on the output layer by selecting a subsetof neurons, we only need to send a subset of embeddings from the embedding table to the client.
It isobvious that this scheme would save the transmission of the largest weight tensor.
Parameters.
In the federated optimization, we use SGD as the client optimizer and Adam as theserver optimizer following FedAdam approach introduced in [22].
For the Adam optimizer in theserver, we vary the epsilon between 10 − 3 and 10 − 4 .
We also perform a grid search on server learningrate set { 10 − 1 , 10 − 1 . 5 , 10 − 2 } and client learning rate set { 10 − 1 , 10 − 1 .
5 , 10 − 2 } .
The training stepsand rounds follow the same setting as [31].
We chose the K and L shown in Section 3.1 from{ 2 , 4 , 6 , 8 } .
We vary random seeds in both model initialization and data loader for a fair comparison.
For the evaluation metrics, we use the accuracy with the out-of-vocabulary, padding, EOS and BOStokens masked.
4.2 Results in Sparse FL
To start with, we would like to evaluate the performance of our hash-based sampling in the outputlayer.
Speciﬁcally, we would like to answer the following question: does hash-based samplingachieves better accuracy than random sampling in the training in LOS with different sparsity?
Herewe conduct an experiment on the Stack Overﬂow dataset, we vary the sparsity level and comparethe hash-based sampling with random sampling.
In Figure 1, we present the evaluation accuracyversus the actual computed parameters.
For each parameter, we repeat the experiment for 3 timesand plot the average accuracy.
As shown in the ﬁgure, if we ﬁx the computed parameter, hash-basedsampling is able to outperform random sampling with better ﬁnal accuracy.
Moreover, we show that,with above 30% of the parameters, the hash-based sampling is able to be less than 0 .
05% full trainingaccuracy.
Noted that the model large vocabulary size can predict more out-of-vocabulary words.
In Figure 1 we do not use a uniﬁed evaluation accuracy across different vocabulary. But we do observe
[Table]
that transformer with 80K vocabulary has better accuracy in the uniﬁed evaluation accuracy.
In thenext section, we would like to show how to use this observation for better on-device training accuracy.
4.3 Model Improvement in Fixed Budget
In this section, we study how to improve the model accuracy with a ﬁx parameter.
Speciﬁcally, wewould like to answer the following question: does the on-device parameters saving in the output layerhelp us improve the model by adding parameters in other parts?
Here we use the parameter size of thefull transformer model described in Section 4.1 as our budget.
We would like to apply sparse trainingin the output layer while increasing the embedding dimension so that we could get closer but notexceed the parameter size budget.
Moreover, with the knowledge from Section 4.2, we only keep 30%of the parameters in the output layer.
In Table 2, we present the results.
For vocabulary size 10K, ifwe use hash-based sampling with 30% compute parameters in the output layer and increase the FFNdimension to 2560, we can outperform the original full model training accuracy in evaluation dataset.
Moreover, if we increase the vocabulary size to 80K, the improvement of hash-based sampling wouldbe enlarged.
If we only select 30% parameters in the output layer using hash-based sampling andincrease the FFN dimension of Transformer to 8192, we could maintain a lower parameter size ondevice. Moreover, we could signiﬁcantly improve the evaluation accuracy to 25 . 2% .
increase the FFN dimension of Transformer to 8192, These experiments validate that our hash-based sparse FLscheme is able to improve the model performance without increasing the on-device parameter size.
Discussion
In this section, we would like to discuss our observations and the potential future directions of thiswork.
Firstly, we observe that for smaller vocabulary sizes, LSH performs marginally better thanrandom sampling.
With our analysis, we observe that LSH does not retrieve large inner productneurons with high recall.
Meanwhile, the exact maximum inner product search (MIPS)on neuronsgives us better accuracy.
In this case, a promising future direction would be the introduction of more
MIPS data structures such as quantization [13] and proximity graphs [39, 28] .
Secondly, we wouldlike to explore the opportunity of further increasing the vocabulary size for the on-device languagemodeling.
Our experimental results have suggested that LSH approach performs better as we increasethe vocabulary.
We would like to investigate how our approach overcomes the on-device hardware limit.
Thirdly, our approach requires communications of hash codes and neurons.
We would like tocombine this approach with more secure and efﬁcient communication schemes [5] for aggregation.
Conclusion
In this paper, we propose a sparse federated learning (FL) scheme using a hash-based adaptivesampling algorithm.
We argue that during the FL training of deep neural networks in large outputspace, we can sample a subset of neurons in the output layer and perform forward and backwardpropagation on these neurons only.
Moreover, we introduce a hash-based adaptive sampling approachin the neuron sampling for FL.
We pre-index the neurons of the output layer in hash tables.
Next,given the input embedding to the output layer, we could look up its near neighbor neurons from hashtables for the sparse training.
Furthermore, we introduce a sparse FL scheme based on this hash-basedsampling approach.
In our scheme, the server takes over the neuron indexing and maintains thehash tables, while the client only maintains a subset of neurons in the last layer through hash tablelookups.
In this way, we show via extensive experiments that we could use around 30% parametersin the last layer and obtain the same ﬁnal accuracy as full parameter training.
We also show that withour approach, we could perform sparse training in the output layer and use the saved parameters toimprove the model capacity in embedding and fully-connected layers.
This design leads us to betteron-device FL accuracy with the same parameter budget.
Appendix
A Visualization of the Hash Table
In this section, we visualize the bucket in the hash table we built for Stack Overﬂow dataset.
Weshowcase the tokens of three buckets from a hash table as below.
Here we set the K = 4 (see
Section 3.1).
java, python, ruby, spring, tomcat, gcc, swift, println, opencv, openssl, activerecord, gdb,clang, webforms [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table] problems ways issues
questions
great
solutions
tutorials
major
bugs
conventions
[Table]
interactions
ﬁndings
[Table] documentations reaction orthogonal
[Table] change, go, deﬁne, modify, ask, perform, manage, determine, accept, reset, combine, main-
[Table]
tain, bring, evaluate, optimize, customize, jump, automate, preserve, terminate, associate, [Table] buy, pip, synchronize, eat, mutate, assemble, recalculate, optimise, dotnet, check-in [Table]
We observe that the ﬁrst bucket is a programming language bucket.
It contains "java", "python" and [Table] related platforms such as "opencv".
In the second bucket, we observe that there are similar tokenssuch as "great", "brilliant" and "useful".
Also the "issue" means similar to "questions" and "bugs".
Inthe third bucket, there is a set of synonyms, such as "change", "modify" and "reset."
To sum up, thehash table is able to group the relative tokens in the same bucket and we can look them up with aquery. [Table]