Efficient Federated Random Subnetwork Training
Abstract
One main challenge in federated learning is the large communication cost of exchanging weight updates from clients to the server at each round.
While prior work has made great progress in compressing the weight updates through gradient compression methods, we propose a radically different approach that does not update the weights at all.
Instead, our method freezes the weights at their initial random values and learns how to sparsify the random network for the best performance.
To this end, the clients collaborate in training a stochastic binary mask to find the optimal sparse random network within the original one.
At the end of the training, the final model is a sparse network with random weights – or a subnetwork inside the dense random network.
We show improvements in accuracy, communication (less than 1 bit per parameter (bpp)), convergence speed, and final model size (less than 1 bpp) over relevant baselines on MNIST, EMNIST, CIFAR-10, and
CIFAR-100 datasets, in the low bitrate regime under various system configurations.
(See (Isik et al., 2022a) for the full version.)
Introduction
Federated learning (FL) is a distributed learning framework where clients collaboratively train a model by performing local training on their data and by sharing their local updates with a server every few iterations, which in turn aggregates the local updates to create a global model, that is then transmitted to the clients for the next round of training.
While being an appealing approach for enabling model training without the need to collect client data at the server, communication of local updates is a significant bottleneck in FL (Kairouz et al., 2021).
This has motivated research in communication-efficient FL strategies (McMahan et al., 2017a) and various gradient compression schemes via sparsification (Lin et al., 2018;
Wang et al., 2018; Barnes et al., 2020; Ozfatura et al., 2021; Isik et al., 2022b), quantization (Alistarh et al., 2017; Wen et al., 2017; Bernstein et al., 2018;
Mitchell et al., 2022), and low-rank approximation (Koneˇcn`y et al., 2016; Vargaftik et al., 2021, 2022; Basat et al., 2022).
In this work, while aiming for communication efficiency in FL, we take a radically different approach from prior work, and propose a strategy that does not require communication of weight updates.
To be more precise, instead of training the weights, the server initializes a dense random network with d weights, denoted by the weight vector w init =
( w init 1 , w init 2 , . . .
, w init d ) , using a random seed SEED , and broadcasts SEED to the clients enabling them to reproduce the same w init locally, 2. both the server and the clients keep the weights frozen at their initial values w init at all times, 3.
clients collaboratively train a probability mask of d parameters θ =
( θ 1 , θ 2 , . .
, θ d ) ∈ [0 , 1] d , 4.
the server samples a binary mask from the trained probability mask and generates a sparse network with random weights – or a subnetwork inside the initial dense random network as follows
[Equation] where Bern ( · ) is the Bernoulli sampling operation and ⊙ the element-wise multiplication.
We call the proposed framework Federated Probabilistic Mask Training (FedPM) and summarize it in
Figure 1. At first glance, it may seem surprising that there exist subnetworks inside randomly initialized networks that could perform well without ever modifying the weight values.
This phenomenon has been explored to some extent in prior work (Zhou et al., 2019; Ramanujan et al., 2020; Pensia et al., 2020; Diffenderfer & Kailkhura, 2020; Aladago & Torresani, 2021) with different strategies for finding the subnetworks.
However, how to find these subnetworks in a FL setting has not attracted much attention so far.
Some exceptions to this are works by Li et al.
; Vallapuram et al.
Mozaffari et al.
(2021), which provide improvements in other FL challenges, such as personalization and poisoning attacks, while not being competitive with existing (dense) compression methods such as QSGD (Alistarh et al., 2017), DRIVE (Vargaftik et al., 2021), and SignSGD (Bernstein et al., 2018) in terms of accuracy under the same communication budget.
In this work, we propose a stochastic way of finding such subnetworks while reaching higher accuracy at a reduced communication cost – less than 1 bit per parameter (bpp).
In addition to the accuracy and communication gains, our framework also provides an efficient representation of the final model post-training by requiring less than 1 bpp to represent (i) the random seed that generates the initial weights w init , and (ii) a sampled binary vector Bern ( θ ) (computed with the trained θ ).
Therefore, the final model enjoys a memory-efficient deployment – a crucial feature for machine learning at power-constrained edge devices.
Another advantage our framework brings is the privacy amplification under some settings, thanks to the stochastic nature of our training strategy.
We show that in the presence of a differential privacy mechanism, such as (Abadi et al., 2016; Agarwal et al., 2021; Andrew et al., 2021), over the probability mask θ , the Bernoulli sampling step amplifies the privacy further under certain conditions.
Our contributions can be summarized as follows:
• We propose a FL framework, in which the clients do not train the model weights, but instead a stochastic binary mask to be used in sparsifying the dense network with random weights.
This differs from the standard training approaches in the literature.
• Our framework provides efficient communication from clients to the server by requiring (less than) 1 bpp per client while yielding faster convergence and higher accuracy than the baselines.
• We propose a Bayesian aggregation strategy at the server side to better deal with partial participation and non-IID data splits.
• The final model (a sparse network with random weights) can be efficiently represented with a random seed and a binary mask which requires (less than) 1 bpp – at least 32 × more efficient storage and communication of the final model with respect to standard FL strategies.
• We demonstrate the efficacy of our strategy on MNIST, EMNSIT, CIFAR-10, and CIFAR100 datasets under both IID and non-IID data splits; and show improvements in accuracy, bitrate, convergence speed, and final model size over relevant baselines, under various system configurations.
Federated Probabilistic Mask Training ( FedPM )
We first describe the simpler version of the FedPM framework in Section 2.1, which provides an unbiased estimation of the mean of the learned probability masks at the server with bounded error .
Next, we propose a modification in our aggregation strategy by exploiting the underlying Bernoulli mechanism in Section 2.2.
This helps boost the performance of FedPM in the cases of non-IID data splits and partial participation of clients.
We then discuss the details of the distribution of the initial weights in Section 2.3, and finally describe the privacy benefits of FedPM in Section 2.4.
Throughout the paper, we use
2.1 FedPM
In this section, we present the general FedPM training pipeline.
First, the server randomly initializes a neural network f w init , parameterized by the weight vector w init = ( w init 1 , w init 2 , . . .
, w init d ) ∈ R d , whose components are sampled IID according to a distribution P w using a randomly generated seed SEED .
The random SEED value is then communicated to all the clients, which can locally sample the same pseudo-random vector w init , which is kept fixed and never modified during training.
The goal for the clients is to collaboratively train a probability mask θ ∈ [0 , 1] d , which indicates the Bernoulli parameters for the global stochastic binary mask M ∼ Bern ( θ ) ∈ { 0 , 1 } d , such that the function f ˙ W maximizes its performance on a given task, where ˙ W = M ⊙ w init .
Specifically, FedPM learns the probabilities for the weights of being active, which are given by the probability mask θ =
( θ 1 , θ 2 , . .
, θ d ) ∈ [0 , 1] d .
To achieve this, at every round t , the server samples a set K t of |K t | =
K participants (out of the total N clients), which individually train their local probability masks θ k,t , k ∈ K t , by using their local datasets D k , each composed of D k = |D k
| samples.
These local masks are then aggregated by the server in a communication-efficient way to estimate the optimal θ .
At test time, at the server, the initial random network f w init is sparsified using the global probability mask θ g,t , following the stochastic approach in Figure 1.
In the following sections, we provide more details on each step of each round.
We give the pseudocode for FedPM in Appendix A.
We reduce the bitrate to less than 1 bit per parameter by using arithmetic coding to encode m k,t . When the frequency of 1 ’s is far from 0 . 5 (which is usually the case with FedPM), the number of bits per parameter to communicate m k,t is less than 1. See Figure 3 for more details.
2.1.1 Local Training of Probability Masks
Upon receiving a global probability mask θ g,t − 1 from the server at the beginning of round t , the client k performs local training and updates the mask via back-propagation.
First, however, we have to guarantee that the updated probability mask satisfies θ k,t ∈ [0 , 1] d .
While this can be achieved with a regularization term or clipping, these approaches may lead to a performance drop due to information loss.
Therefore, similarly to the work of Zhou et al.
(2019), we introduce another mask, called score mask s =
( s 1 , s 2 , . .
, s d ) ∈ R d , that has unbounded support and that can be used to generate the probability masks through the one-to-one sigmoid function by setting θ =
Sigmoid ( s ) .
Then, the procedure for local training of the probability mask at round t is as follows (here, the steps from Step 2 to 4 describe one local iteration, which is repeated a number τ of times as standard in FL (McMahan et al., 2017a)):
The server sends the global probability mask θ g,t − 1 to K chosen clients, and the clients set s k,t =
Sigmoid − 1 ( θ g,t − 1 ) , where Sigmoid − 1 ( · ) is the inverse of the sigmoid function.
Then, the clients generate a binary mask by first transforming back θ k,t =
Sigmoid ( s k,t ) , and then sampling a binary mask M k,t from θ k,t as shown in Figure 1: m k,t ∼ Bern ( θ k,t ) .
The sampled binary mask then sparsifies the initial weight vector w init : ˙
w k,t
= m k,t ⊙ w init . 4. ˙
w k,t is then used for forward pass, and the loss L ( f ˙ w k,t , D k ) on the local task is backpropagated to update the score mask as s k,t = s k,t − η ∇L ( f ˙ w k,t , D k ) ( η is the local learning rate).
All the local operations from Step 2 to Step 4 are differentiable, except for the Bernoulli sampling.
To backpropagate the gradients through the sampling operation, we use the first-order gradient of the Bernoulli function, which is simply equal to the probability mask θ k,t .
2.1.2 Communication Strategy
Once the local training at round t is completed, the server needs to distill the global probability mask θ g,t , for example, by taking the empirical average of the local probability masks ¯ θ g,t =
K (cid:80) k ∈K t θ k,t from the clients.
However, since we aim for communication efficiency, the clients do not send their local probability masks directly.
Instead, they communicate a stochastic binary sample
M k,t from their probability masks sampled as m k,t ∼ Bern ( θ k,t ) , and then the server estimates the global aggregate ¯ θ g,t as ˆ¯ θ g,t = 1 K (cid:80) k
∈K t m k,t .
This distributed mean estimation problem with communication constraints is summarized in Figure 2.
Our estimator ˆ¯ θ g,t = 1 K (cid:80) k ∈K t m k,t is an unbiased estimate of the true aggregate, in that
[Equation]
Moreover, the estimation error is upper bounded as (the proof is given in Appendix B)
[Equation]
Since each client communicates a stochastic binary mask M k,t , 1 bpp is the worst case bitrate for FedPM.
We can further reduce the bitrate to less than 1 bit by using arithmetic coding (Rissanen &
Langdon, 1979) or universal coding (Krichevsky & Trofimov, 1981; Barron et al., 1998) to encode m k,t , and achieve the empirical entropy since d is large.
This gives us smaller bitrates whenever the frequency of 1’s in m k,t is far from 0.5 – which is usually the case for our method (see Figure 3 and Appendix E.1 for results).
We note that, with a deterministic mask training approach as in FedMask (Li et al., 2021), arithmetic coding of m k,t s does not provide any further gain in bitrate, as we have empirically observed that the frequency of 1’s is always around 0.5 (see Figure 3 and Appendix E.1) – here we apply arithmetic coding for FedMask to improve our baseline although it was not proposed in the original paper.
Moreover, FedMask (Li et al., 2021) and HideNSeek (Vallapuram et al., 2022) do not enjoy the guarantees we have as their estimator (i) is not unbiased and (ii) does not have an upper bound on the estimation error due to hard thresholding (Li et al., 2021) and sign operations (Vallapuram et al., 2022).
This is another benefit of our stochastic sampling approach.
2.2 FedPM with Bayesian Aggregation
Another important aspect that differentiates our work from existing masking methods such as FedMask (Li et al., 2021) and HideNSeek (Vallapuram et al., 2022) is the Bayesian aggregation strategy, which exploits the underlying stochastic mask to synthesize a global model, boosting the performance in heterogeneous scenarios, e.g., when local client data are not sampled from the same distribution.
Given the probabilistic interpretation of the FedPM mask’s values, at the server side we further model the probability mask θ g,t with a Beta distribution Beta ( α g,t , β g,t ) , parameterized by the round-dependent parameters α g,t and β g,t , which are initialized to α g, 0 = β g, 0 = λ 0 .
At the beginning of the training process, there is no prior knowledge indicating which network weight should be more important than the others, and so each entry in the probability mask is uniformly distributed in [0 , 1] – which is the prior distribution.
Consequently, the clients’ local binary masks M k,t s are the data the server uses to update its belief on each weight score, and so the aggregation strategy corresponds now to a posterior update .
Specifically, given the conjugate relation between the Beta-Bernoulli distributions, the new posteriors are still Beta distributions with parameters α g,t
= α g,t − 1 + M agg ,t and β g,t = β
[Equation]
where M agg ,t = (cid:80) k ∈K t M k,t , and 1 is the d -dimensional vector containing all ones.
Then, the server broadcasts to the clients the mode of the Bernoulli distributions, as suggested by Ferreira et al.
(2021), i.e., [Equation] where the division operation is applied element-wise.
However, to obtain the best performance out of this method, the Beta parameters should be re-initialized to their original values λ 0 with some regularity.
Notice that if λ 0 = 1 , and if α and β are re-initialized at the beginning of each round, the method is equivalent to the aggregation strategy detailed in Section 2.1.2.
2.3 Weight Distribution
As mentioned in Section 2.1, the fixed weight vector w init is initialized by sampling from the distribution P w using the randomly generated SEED .
We note that the choice of this distribution impacts two important aspects of FedPM: (i) the values of w init highly influence the final accuracy achieved by the model, as they represent the building blocks to extract a subnetwork f ˙ w (see Figure 1), which should be rich enough to solve the learning task, and (ii) the size of the sample space of P w affects the number of bits needed to store the model during the inference process (this is different from the 1 bpp model storage when the model is not in use). Regarding (i), as also proposed in Ramanujan et al. , we sample weights from a uniform distribution, whose domain is {− σ, + σ } , where σ is the standard deviation of the Kaiming Normal distribution (He et al., 2015).
In this way, we control the variance of the neurons’ output to be ∼ 1 , which avoids the vanishing or the explosion of activation values. Previous experiments in (Zhou et al., 2019; Ramanujan et al., 2020) also demonstrate the superior performance achieved by binary weights distributions when compared to standard continuous counterparts, e.g., Gaussian. Regarding (ii), even if knowing the value of SEED is enough to perfectly reconstruct the vector w init , one would have to generate the entire vector at every inference step.
Consequently, to achieve fast inference, the actual values of the weights need to be stored in the memory of the devices during the inference process .
Fortunately, our initialization allows for efficient storage even during inference since (after reconstructing w final using SEED and m final ∈ { 0 , 1 } d ) we only need to indicate whether the weight values in w final are − σ , 0 , or + σ , with a ternary representation that can be efficiently deployed on hardware (Alemdar et al., 2017).
2.4 Privacy
Privacy is another challenge in FL as the model updates (in our case, M k,t s) may leak information about the client data.
Differential privacy (DP) guarantees that the probability of an outcome of an algorithm that runs on client data does not change much by a single client’s data (see Appendix C for the formal definition).
This is typically ensured via injecting noise to a function of the client data at a particular step in the algorithm with some utility loss in the application.
While there have been many DP strategies developed for FL and deep learning (Abadi et al., 2016; McMahan et al., 2017b; Agarwal et al., 2021; Andrew et al., 2021), these strategies typically suffer from severe performance degradation due to noise injection.
To make DP practical, researchers have explored certain randomization mechanisms that amplify the privacy guarantee.
When these mechanisms are part of the FL framework, such as sampling (data (Balle et al., 2018;
Wang et al., 2019) or device (Balle et al., 2020; Girgis et al., 2021; Hasircioglu & Gunduz, 2022)) and shuffling (Erlingsson et al., 2019; Feldman et al., 2022), the amplification comes for free.
This is helpful because the overall process can meet a stronger privacy guarantee without increasing the noise level.
FedPM promises one such amplification due to the stochastic Bernoulli sampling step.
In particular, Imola & Chaudhuri (2021) have shown that when a sample M ∈ { 0 , 1 } d from an already privatized vector θ ∈
[ c, 1 − c ] d , where 0 < c < 0 .
5 , is released to a third party (instead of θ itself), the privacy is amplified under some conditions.
More precisely, when there is an ( α, ϵ ) -Rényi
Differential Privacy mechanism (Mironov, 2017) that privatizes θ ∈ [ c, 1 − c ] d , releasing a sample from Bern ( θ ) yields an improved privacy budget (the smaller ϵ , the better the privacy): ϵ amp ≤ min { ϵ, d · r α ( c ) } .
Here, r α ( p ) is the binary symmetric Rényi divergence function defined as r α ( p ) =
α − 1 log (cid:0) p α (1 − p ) 1 − α + (1 − p )
α p 1 − α (cid:1) .
Notice that FedPM already involves this Bernoulli sampling step in the communication protocol and in the forward pass m k,t ∼ Bern ( θ k,t ) .
However, the d term in the upper bound limits the amplification for large model sizes.
We believe it is worth exploring a tighter upper bound on ϵ amp to enjoy privacy amplification in FedPM with practical models.
Nonetheless, in Appendix C, we demonstrate the impact of this amplification on a distributed mean estimation problem, described in
Figure 2, where the goal is to estimate the true mean of the probability masks ¯ θ = 1 K (cid:80)
Kk
=1 θ k under communication and privacy constraints.
We also provide a bias correction mechanism, specific to our scheme in Figure 4 in Appendix C, that mitigates the bias due to the DP mechanism and reduces the estimation error.
Experiments
In this section, we empirically show the performance of FedPM in terms of accuracy, bitrate, converge speed, and the final model size.
We consider four datasets: CIFAR-10 with 10 classes, CIFAR-
(Krizhevsky et al., 2009) with 100 classes, MNIST (Deng, 2012) with 10 classes, and EMNIST (Cohen et al., 2017) with 47 classes.
For CIFAR-100, we use a 10-layer convolutional network (CNN) CONV - 10 ; for CIFAR-10, a 6-layer CNN CONV - 6 ; and for MNIST and EMNIST, a 4-layer CNN
CONV - 4 .
A detailed description of the architectures can be found in Appendix D.
We first compare
FedPM with SignSGD (Bernstein et al., 2018), TernGrad (Wen et al., 2017), QSGD (Alistarh et al., 2017), DRIVE (Vargaftik et al., 2021), EDEN (Vargaftik et al., 2022), and FedMask (Li et al., 2021) on IID data split and full client participation in Section 3.1.
We then extend our experiments to non-IID data splits and partial participation in Section 3.2.
Clients perform 3 local epochs in all experiments.
We provide additional details on the experimental setup in Appendix D. We present results averaged over 3 runs.
3.1 IID Data Split and Full Participation ( K = N )
In this section, we focus on IID data distribution and the case when all the clients participate in the training at each round.
We set the number of clients to N = K = 10 .
We report the estimated bitrate for the arithmetic code that uses the empirical frequency of the symbols (for our method FedPM, this corresponds to the frequency of 1’s in m k,t ) – which is equal to the empirical entropy for blocklength d as large as the model size.
In Figure 3, we compare the accuracy, bitrate, and convergence speed of FedPM with relevant baselines.
As can be seen in the figure, FedPM converges to the highest accuracy on all four datasets.
DRIVE, EDEN, and QSGD (they mostly overlap in the accuracy plots) seem to be the three baselines that perform the best after FedPM; however, their convergence speed is significantly lower than FedPM.
In terms of convergence speed, FedMask is the fastest among the baselines – in fact, at the beginning of the training, FedMask is faster than FedPM as well.
However, its final accuracy is lower than the others.
In terms of bitrate, SignSGD and FedMask consistently spend 1 bpp, which is the default number when a binary mask or sign mask is communicated.
This means binary values (1’s and 0’s) are almost equally distributed in their masks, which prevents them from enjoying additional bitrate gains.
Across all experiments, TernGrad has the highest bitrate.
We would like to leave a note about the bitrate of QSGD.
Unlike other baselines, including our work, QSGD can go down to very low bitrates by adjusting the number of levels in quantization.
We have observed that in the extreme quantization case, QSGD underperforms FedPM. Then, we have decided to increase the number of quantization levels in QSGD to see if it improves the accuracy.
However, as can be seen from the plots, even with bitrate larger than 1, QSGD still underperforms FedPM.
The only two baselines that challenge FedPM in terms of bitrate are DRIVE and EDEN.
While FedPM has lower bitrates on CIFAR-10 and EMNIST; DRIVE and EDEN have better bitrates on CIFAR-100 and MNIST.
However, the accuracy of DRIVE and EDEN on these datasets (specifically CIFAR-100) is significantly lower than that of FedPM, with slower convergence.
As for the final model size, FedPM needs only 0.8 bpp for the CONV - 6 model trained on CIFAR-10, 0.85 bpp for the CONV - 10 model trained on CIFAR-100, 0.96 bpp for the CONV - 4 model trained on
MNIST, and 0.83 bpp for the CONV - 4 model trained on EMNIST.
On the other hand, other baselines that train a dense model, namely SignSGD, TernGrad, QSGD, DRIVE, and EDEN, would need to represent each weight with their full precision value, i.e., 32 bpp.
This implies that FedPM provides around 38 .
6 × improvement in the storage or the communication of the final model.
Since FedMask also trains a sparse model, it enjoys a similar gain in the final model size requiring 1 bpp across all the models.
Due to the stochastic masking procedure and uneven distribution of 1 ’s and 0 ’s in the binary masks, FedPM has up to 0.17 bpp improvement over the deterministic procedure in FedMask, which adds up to a large gain due to the huge model size.
We would like to highlight that while some of our baselines, such as FedMask and TernGrad, have a visibly high variance in accuracy, FedPM shows stable training behavior across all experiments.
3.2 Non-IID Data Split and Partial Participation ( K < N )
This section considers more realistic scenarios, in which the local clients’ datasets are generated from slightly different data distributions.
We focus on CIFAR-10 with CONV - 6 , and we compare FedPM against (i) the most promising baselines, which, based on the results of Section 3.1, are
DRIVE, EDEN, and QSGD, and (ii) FedMask, as it is the only sparse baseline. To choose the size of each dataset |D n | = D n , for each client n ∈ { 1 , . . . , N } , an integer j n is sampled uniformly from { 10 , 11 , . .
Then, a coefficient p n = j n (cid:80) n j j is computed, which represents the size of the local dataset D n as a fraction of the size of the full dataset, i.e., the training set of CIFAR-10. In this way, highly unbalanced datasets can be generated from the central one. Moreover, since the task is a classification problem, we impose a maximum number of different labels, or classes, c max , that one client can see.
Consequently, clients need cooperation to learn the statistics of other classes’ distributions, as the test dataset contains samples from all classes. In addition, partial participation is also considered, meaning that at each round, the server uniformly samples a fraction ρ = KN of the clients to participate in the training round. This is motivated in real-world scenarios by the scarcity of physical communication network resources, which may limit the availability of part of the clients during one round. The maximum number of classes per local dataset is set to c max ∈ { 2 , 4 } , and the participation ratio is set to ρ ∈ { 0 . 1 , 0 . 5 , 1 } .
For ρ = 1 and ρ = 0 . 5 , the total number of clients is set to N = 10 (and so K is equal to 10 and 5 , respectively).
For ρ = 0 . 1 , we set N = 50 (and so
K = 5 ), which is the worst scenario among all combinations, given the small amount of information the server can collect at the end of each round.
When ρ = 1 , for the FedPM algorithm, we keep the same aggregation strategy exposed in Section 2.1.2 and Figure 2; and we switch to the Bayesian aggregation method (see Section 2.2) when there is partial participation, i.e., when ρ < 1 .
Indeed, applying the Bayesian aggregation method is revealed to be crucial for achieving good accuracy when ρ < 1 and data are non-IID, obtaining a large gain with respect to the simpler version in Section 2.1.2, which resets the Beta priors at each round (or takes the average of the samples, as explained in Section 2.2).
We adopt a simple heuristic schedule to reset the priors: Reset every 3 rounds when ρ = 0 .
5 , and every 10 rounds when ρ = 0 .
As expected, the smaller the ratio ρ , the larger the number of rounds we should wait before resetting the priors in order to collect more information from a much more diverse pool of clients.
Table 1 reports the results when c max = 4 and c max = 2 .
As we can see, FedPM outperforms all the baselines in every configuration, as the Bayesian aggregation allows the central node to collect more data before resetting the priors, which is important when clients’ data distributions are heterogeneous.
This strategy can be seen as the FedPM counterpart of decreasing the learning rate (which we applied in the other dense compression-based baselines, like DRIVE, EDEN, and QSGD).
It is seen from
Table 1 that FedMask (Li et al., 2021) is struggling in the non-IID case, as applying a hard threshold on the scores to binarize the mask does not provide a proper way to implement multiple-rounds aggregation, emphasizing the benefit of the stochastic process in FedPM.
It is interesting to notice that, especially when c max = 4 , the lower the value of ρ , the larger the gap between FedPM and the [Table] baselines, corroborating the fact that the Bayesian strategy can better deal with partial participation.
Analysis of the communication bitrate is provided in Appendix. E.1.
Conclusion
In this work, we introduced Federated Probabilistic Mask Training ( FedPM ) – a communicationefficient FL strategy.
FedPM relies on the idea of finding a sparse network in a randomly initialized dense network, which is then sparsified by a collaboratively trained stochastic binary mask.
In addition to reducing the communication cost to less than 1 bit per parameter (bpp), FedPM also reaches higher accuracy with faster convergence than the relevant baselines, and can potentially amplify privacy while additionally outputting a compressed final model with a size less than 1 bpp.
Throughout the manuscript, we highlighted the advantages of having a stochastic mask training approach rather than a deterministic one in terms of accuracy, bitrate, and privacy.
For instance, the proposed Bayesian aggregation strategy boosts the performance of FedPM in the non-IID data split and/or the partial client participation case by using the prior knowledge from previous rounds during aggregation rather than hard-replacing the previous global probability mask with the new one.
Acknowledgement
The authors would like to thank Zachary Charles, Mahdi Haghifam, Peter Kairouz and Nicole
Mitchell for inspiring discussions.
This work was supported in part by a Sony Stanford Graduate
Fellowship and a Meta research award.
A FedPM Algorithm
We provide the pseudocode for FedPM in Algorithms 1 and 2.
In Algorithm 2, the prior resetting scheduling policy is controlled by the procedure ResPrior( t ) , which may depend on quantities other than the round number t , such as loss.
Algorithm 1 FedPM.
Hyperparameters: learning rate η , minibatch size B , number of local iterations τ .
Inputs: local datasets D i , i = 1 , . . . , N
Output: random seed SEED and binary mask parameters m k,T
At the server, initialize a random network with weight vector w init
∈ R d using a random seed SEED , and broadcast it to the clients. At the server, initialize the random score vector s g, 0 ∈ R d , and compute θ g, 0 ← Sigmoid ( s g, 0 ) .
At the server, initialize Beta priors α g,
= β g, 0 = λ 0 .
for t = 1 , . . .
, T do
Sample a subset K t ⊂ { 1 , . . . , N } of |K t | =
K clients without replacement.
On Client Nodes: for k ∈ K t do
Receive θ g,t − 1 from the server and set s k,t =
Sigmoid − 1 ( θ g,t − 1 ) .
for l = 1 , . . .
, τ do θ k,t ← Sigmoid ( s k,t )
Sample binary mask m k,t ∼ Bern ( θ k,t ) .
˙ w k,t ← m k,t ⊙ w init grad s k,t ← 1 B (cid:80) Bb =1 ∇ ℓ ( ˙ w k,t ; B kj ) ; {B kj } Bj =1 is uniformly chosen from D k s k,t ← s k,t − η · grad s k,t end for θ k,t ← Sigmoid ( s k,t )
Sample a binary mask m k,t ∼ Bern ( θ k,t ) .
Send the arithmetic coded binary mask m k,t to the server.
end for
On the Server Node:
Receive m k,t ’s from K client nodes.
θ g,t = BayesAgg( { m k,t See Algorithm Algorithm 2 BayesAgg.
Output global probability mask θ g,t if ResPriors( t ) then [Equation]
end if
Compute
[Equation]
Return g, [Equation]
B Proof of the Upper Bound on the Estimation Error
We now provide proof of the upper bound on the estimation error in Eq. 2.
Recall that our true mean is ¯ θ g,t = 1 K (cid:80) k ∈K t θ k,t , whereas our estimate is ¯ θ g,t = 1 K (cid:80) k ∈K t m k,t , where m k,t ∼ Bern ( θ k,t ) .
Then we can compute the error as
[Equation]
t E
[Equation]
k
[Equation]
ti
[Equation]
Bern ( θ k
[Equation]
t
[Equation]
) (cid:104) (
[Equation]
(cid:105) (9) =
[Equation]
(cid:88) k
[Equation]
(cid:16)
[Equation]
(10) =
[Equation]
(cid:16)
[Equation]
(11) ≤
[Equation]
[Equation]
From (5) to (6), we use the definition of ˆ¯ θ g,ti = 1 K (cid:80) Kk
{+=
+}
m k,ti and ¯ θ g,ti = 1 K (cid:80)
Kk
=1 θ k,ti .
From (7) to (8), we use the fact that E M k,ti ∼ Bern ( θ k,ti ) ∀ k ∈K t
[ M k,ti − θ k,ti ] = 0 ; and M k,ti − θ k,ti and M l,ti − θ l,ti are independent for l ̸ = k ∈ [ K ] .
Finally, the inequality in (8) follows from θ k,ti ∈ [0 , 1] for all k ∈
[ K ] .
C Privacy Amplification and Bias Correction
We first revisit the definitions of differential privacy (Dwork et al., 2006), Rényi divergence, and
Rényi differential privacy (Mironov, 2017).
Definition 1.
[Adjacent Datasets] Two datasets D, D
′ ∈
D are called adjacent if they differ in at most one data sample.
Definition 2.
[ ( ϵ, δ ) -DP] A randomized mechanism f : D → R offers ( ϵ, δ ) -differential privacy if for any adjacent D, D ′ ∈ D and S ⊂
[Equation]
Definition 3.
[Rényi Divergence]
For two probability distributions P and Q defined over R , the
Rényi divergence of order α > 1 is
[Equation] Definition 4. [ ( α, ϵ ) -RDP] A randomized mechanism f : D → R offers ϵ -Rényi differential privacy of order α (or in short ( α, ϵ ) -RDP) if for any adjacent D, D ′ ∈ D , it holds that
[Equation]
Now, recall our discussion in Section 2.4 that we have an ( α, ϵ ) -RDP algorithm f that outputs privatized θ k ∈ [ c, 1 − c ] d using local client data D k . As summarized in Figure 4, we are interested in what happens when instead of releasing θ k
= f ( D k ) , the client k releases a Bernoulli sample from it: m k ∈ { 0 , 1 } d ∼ Bern ( θ k ) .
We already explained the advantages in terms of communication bitrate, estimation error, unbiasedness throughout the manuscript; however, this approach also amplifies the privacy guarantees, meaning that it makes the overall privacy budget smaller ϵ amp ≤ ϵ . Quantitatively, Imola & Chaudhuri (2021) showed that after the Bernoulli sampling, the privacy budget of the overall process is
[Equation]
where r α ( · ) is the Rényi divergence of the binary symmetric function.
More precisely, consider P, Q random variables with support on { x 1 , x 2 } ⊂ Θ and let p =
Pr [ P = x 1 ] , 1 − p =
Pr ( Q = x 1 ) .
Then the Rényi divergence is defined as
[Equation]
This implies that the stochastic Bernoulli sampling step of FedPM improves the privacy guarantee without changing the privacy mechanism – e.g. without increasing the injected noise level. We demonstrate this through a distributed mean estimation problem given in Figure4.
We consider a simple differential privacy (DP) strategy, where the probability masks θ k ∈ [ c, 1 − c ] d are a function of client data D k ; and are first corrupted by Gaussian noise, and then clipped to the range [ c, 1 − c ] d .
Our goal is, as before, to estimate the true mean ¯ θ = 1 K (cid:80) k ∈K t θ k by averaging the sampled binary masks, i.e., ˆ¯ θ = 1 K (cid:80) k ∈K t m k .
Differently from our previous experiments, we have privacy constraints now, meaning that we want to guarantee ( ϵ, δ ) -DP by injecting a Gaussian noise with variance σ 2 = 2ln(1 .
25 /δ )
∆ 22 ϵwith a small ϵ , where δ ≈ 1 N 2 and ∆ 2 is the ℓ 2 -sensitivity of the probability masks (in our case ∆ 2 = (1 − 2 c ) √ d ).
We transfer the above amplification results in
RDP to DP using the well-known relation: Remark C.1. Mironov (2017) showed that if f is an ( α, ϵ ) -RDP mechanism, it also satisfies ( ϵ + log1 /δ
α − 1 , δ ) -DP for any 0 < δ < 1 .
Since clipping after the noise addition step would lead to bias in the estimated mean, we work out a bias correction mechanism.
We denote with θ one general parameter at client k for one parameter, with ˜ θ its noisy version, and with ˆ θ = clip (˜ θ ) its clipped version.
Specifically, if ˜ θ
= θ + η is the noisy version of the parameter, where η ∼ N (0 , σ 2 ) , then clip
[Equation]
We now compute E (cid:104) ˆ M (cid:105) , where ˆ M ∼ Bern (ˆ θ ) , to analyze the bias E (cid:104) ˆ M (cid:105) − E [ M ] = E (cid:104)
ˆ M (cid:105) − θ , where M ∼ Bern ( θ ) .
First of all, notice that
[Equation]
And we now compute the mean of the clipped parameter
[Equation]
where Φ σ ( · ) is the cumulative distribution function of a Gaussian random variable with standard deviation σ , and zero mean.
We use this relation to correct the bias in ˆ¯ θ .
We conduct our experiments with N = 100 clients, each having independent probability masks with dimension d = 5 and range [0 . 2 , 0 . 8] , i.e., θ ∈ [0 . 2 , 0 . 8] 5 .
Figure 5 shows the estimation error
|| ˆ¯ θ g,t − ¯ θ g,t || 22 under no noise injection case (i.e. no DP) with the black line.
Recall that we want to reach a smaller estimation error and smaller ϵ (i.e., a stronger privacy guarantee). The red curve corresponds to the ϵ vs. estimation error behavior if Bernoulli sampling did not amplify the privacy.
The blue curve shows the amplified ϵ (i.e. ϵ amp ≤ ϵ ) vs. estimation error behavior, and it overlaps with the red curve for ϵ values smaller than d · r α ( c ) = 8 .
96 , where there is no privacy amplification, i.e., ϵ amp
= ϵ . However, notice that the blue line never reaches ϵ ’s higher than this value due to amplification, while enjoying smaller estimation errors that the red curve can only achieve with very large ϵ .
This shows the promise of FedPM in having a better privacy-accuracy performance than most baselines that do not have amplification.
Finally, the green curve shows that bias correction improves this performance further even with ϵ < d · r α ( c ) = 8 . 96 by achieving lower estimation errors with the same ϵ .
D Additional Experimental Details
In Table 2, we provide the architectures for all the models used in our experiments.
Clients performed 3 local epochs with a batch size of 128 and a local learning rate of 0.1 in all the experiments.
Notice that there is no server learning rate in FedPM; instead, we tune the prior resetting schedule in Bayesian aggregation for the experiments in Section 3.2.
We conducted our experiments on NVIDIA Titan X
GPUs on an internal cluster server, using 1 GPU per one run.
In the non-IID and partial participation experiments in Section 3.2, to distill the final model, we may apply both stochastic sampling, as during training, or a hard-threshold method, similar to the one [Table] adopted in FedMask (Li et al., 2021).
In the latter [Table] a binary mask coefficient m i is set to 1 if θ i > α ths , and 0 otherwise.
For all experiments but one, when α ths ∈ [0 . 4 , 0 . 6] , the thresholding test accuracy is always higher than the sampling method
[Table]
and so we use the threshold method.
However, in the extreme case c max = 2 and ρ = 0 . 1 , the optimal values for α max were in [0 . 2 , 0 . 4] and [0 . 6 , 0 . 8] in all experiments, probably due to the high randomness given by the highly heterogeneous scenario.
Consequently, for the last experiment, we just adopt the stochastic sampling strategy to evaluate the model, as further optimizing the α ths means adapting to the test dataset, which may corrupt the ability of the model to generalize.
E Additional Experimental Results E.1 Bitrate Considerations on non-iid Data
We now report the communication bitrate considerations on the non-IID data split experiments described in Section 3.2.
Table 3 reports the average bitrate needed by different algorithms over the whole training process when c max = 4 and c max = 2 .
By simply multiplying the obtained average bitrate by the total number of rounds t max = 200 , we obtain the total number of bits one element in the global probability mask needs to converge to its final value, indicating the total amount of information communicated during the training process.
We first observe that both DRIVE and EDEN consume almost the same amount of bits no matter the system configuration and round number (very small variance), and it is instead model dependent (see Figure 3).
On the contrary, FedPM and QSGD report higher bitrate variability, as it depends on both the training phase and system setting.
As already observed in Section 3.1, FedMask balances almost uniformly the binary updates, leading to a bitrate that is basically fixed to 1 .
For both c max = 4 and c max = 2 , FedPM yields the smallest bitrate when ρ = 1 , whereas for the other scenarios, EDEN and DRIVE are slightly more efficient.
We argue that this is motivated by the fact that, as the learning task becomes harder due to the high system heterogeneity, all the models struggle to converge to good and stable solutions, which means that FedPM is still uncertain about the weights’ importance probabilities θ , setting many of them close to 0 .
However, we think that this may be a useful feature of FedPM to quantify its internal uncertainty, which we will further analyze.
To conclude the analysis, we also report the FedPM bpp for the final model, which is an indication of the average number of bits needed per one parameter of the model.
In the case of c max = 4 , the final model sizes are 0 .
79 bpp, 0 . 834 bpp, and 0 .
99 bpp, when ρ = { 0 . 1 , 0 . 5 , 1 } , respectively.
When c max = 2 , the final model sizes are 0 . 8 bpp, 0 .
817 bpp, and 0 . 992 bpp.
Consequently, at the end of the training process, FedPM remains the most efficient option, as already observed in Section 3.1.
[Table]