FEDDROP: TRAJECTORY-WEIGHTED DROPOUT
Year: 2021

Abstract:
Federated learning (FL) enables edge clients to train collaboratively while preserving individual’s data privacy. As clients do not inherently share identical data distributions, they may disagree in the direction of parameter updates, resulting in high compute and communication costs in comparison to centralized learning. Recent advances in FL focus on reducing data transmission during training; yet they neglected the increase of computational cost that dwarfs the merit of reduced communication. To this end, we propose FedDrop, which introduces channel-wise weighted dropout layers between convolutions to accelerate training while minimizing their impact on convergence. Empirical results show that FedDrop can drastically reduce the amount of FLOPs required for training with a small increase in communication, and push the Pareto frontier of communication/computation trade-off further than competing FL algorithms.

heading: 1 INTRODUCTION
text: In the light of the importance of personal data and the recent strict privacy regulations, e.g. the General Data Protection Regulation (GDPR) of the European Union (Voigt & Von dem Bussche, 2017; Wolters, 2017; Politou et al., 2018), there is now a great amount of risk, responsibility (Edwards et al., 2016; Culnan & Williams, 2009) and technical challenges for securing private data centrally (Sun et al., 2014); it is often impractical to upload, store and use data on central servers. To this end, federated learning (FL) (McMahan et al., 2017; Li et al., 2019a) enables multiple edge compute devices to learn a global shared model collaboratively in a communication-efficient way without collecting their local training data. When compared against naïve decentralized SGD, Federated averaging (FedAvg) (McMahan et al., 2017) and subsequent FL algorithms (Li et al., 2020; Karimireddy et al., 2020) reduce the burden of data transmission by many orders of magnitude.
While these communication-efficient methods can notably alleviate the difficulties of using FL in scenarios with limited bandwidth or data-quota, they, however, entail a drastic increase in computation cost, which has rarely been addressed by previous literature on FL. It has been argued that communication is several orders of magnitude more expensive than computation on edge devices (Li et al., 2019a; Huang et al., 2013). Yet existing FL methods optimize for communication so aggressively that an iPhone 12 Pro running as a FedAvg client in our test1 would spend at least 16 minutes on training, before even transmitting 180 MB of data, which only takes fraction of a second under the modern 5G infrastructure. The savings in communication thus are completely dwarfed by the expensive computation.
In a FL setup, models trained by the clients may disagree on the gradient direction to update the same neurons if the user data distributions are non-IID. In Figure 1, we demonstrate this effect with an initial round of FedAvg training. It turns out that clients sharing similar data distributions tend to produce similar update trajectories to the same channel, and different distributions observe disparate trajectories. From this example, it can be observed that the consequence of non-IID data is two-fold. First, a naïve averaging of client parameters may cause these conflicting signals to cancel out each other, resulting in a slow convergence of the global model. Second, neurons in a layer tend to learn distinct features, and yet they cannot learn meaningfully if these features are absent in the client’s training data. Since neuron training is heavily dependent on the client’s local data, it presents us an
1See Appendix F.6 for the setup.
opportunity: can we concentrate training effort to neurons that are correlated to the current data distribution of the client, while paying less attention to neurons that are less relevant to the client? To leverage this, we introduce FedDrop, which introduces SyncDrop layers to structurally sparsify client models, and on the server-side an inter-client trajectory-based optimization of the dropout probabilities used by SyncDrop to speed up model training. FedDrop brings two-fold advantages: dropped neurons can simply be skipped, and thus reduce the training FLOPs required per step; and dropout probabilities of each neuron can be tuned individually to minimize the impact on the global averaged model updates to assist convergence. Overall, as evinced by our experiments, FedDrop enables a much improved communication/computation trade-off relationship than traditional FL approaches.
Our contributions in this paper are as follows:
• We present SyncDrop layers, which are synchronous structural dropouts with adaptive keep probabilities, to reduce the computational cost required by FL.
• With SyncDrop, we formally derive the FedDrop objective that adjusts each neuron’s dropout probability to improve convergence by minimizing the disparities among inter-client update trajectories. We further introduce FLOPs-based constraints to enforce sparsity per FL round, allowing the trade-off between FLOPs and communication to be tuned easily.
• Empirical results reveal that the combined method, FedDrop, attains a substantially better communication/computation trade-off in comparison to other FL methods.

heading: 2 RELATED WORK
text: Federated learning. Distributed machine learning has a long history of progress and success (PeteiroBarral & Guijarro-Berdiñas, 2013; Li et al., 2014), yet it mainly focuses on training with IID data. The Federated Learning (FL) paradigm and the Federated Averaging algorithm (FedAvg) initially introduced by McMahan et al. (2017) allow clients to train collaboratively without sharing the private data in a communication-efficient manner, To further tackle data heterogeneity, FedProx (Li et al., 2020) introduces new regularizations, and SCAFFOLD (Karimireddy et al., 2020) presents control variates to account for client drifts and reduce the inter-client variance. While being effective at reducing communication, the above methods neglected the computational costs associated with the training process.
Computation vs. communication during training. There are a few precursory methods that focus on the joint optimization of computation and communication costs during training. Caldas et al. (2018a) introduced federated dropout, which prunes parameters following a uniform random distribution, whereas PruneFL (Jiang et al., 2019) proposes a greedy gradient-magnitude-based unstructured pruning. In each FL round, both methods produce a shared pruned model with fine-grained sparsity for all clients. Such unstructured sparsity is difficult to utilize for training acceleration, and a shared global model cannot exploit the data distribution of individual clients. Adaptive federated dropout (Bouacida et al., 2020) partially addresses the latter issue by allowing each client to select a sub-model to join training. FjORD (Horvath et al., 2021) introduces ordered dropout to tackle the problem of system heterogeneity in FL with sub-model training that dynamically changes model capabilities using uniform sampling. FedGKT (He et al., 2020) transmits shallow layer activations to offload training of subsequent layers to the server. However, the privacy implications were not well explored, and the trained models were substantially larger, which may limit their inference speed on edge devices. It is noteworthy that FedDrop differs from these approaches as it takes into account the non-IID client data distributions in a FL setting, and computes inter-client channel selection decisions that would minimize the impact on convergence.
Dropout algorithms in centralized training. Dropout (Hinton et al., 2012; Srivastava et al., 2014) may improve neural network training by reducing the overfitting effect and allow it to generalize better. It is done by randomly setting parts of the connections or weights to zero in a network, and scaling the remaining values accordingly. Since its advent, there have been an increasing interest in applying dropout in a structural fashion (Huang et al., 2015; Ghiasi et al., 2018; Hou & Wang, 2019). Nonetheless, the above methods missed the opportunity to explore the implications of a structural dropout on the FLOPs consumed by training or inference.
Structural pruning. A closely related topic is structural neuron pruning/selection for accelerated inference. These work propose new ways to extract a smaller accurate network from the original (He et al., 2018; Wu et al., 2019; Li et al., 2019b; Herrmann et al., 2020), and can even do so dynamically with an input-dependent policy (Gao et al., 2019; Hua et al., 2019; Wang et al., 2020c).
3 THE FEDDROP METHOD
3.1 HIGH-LEVEL OVERVIEW
FedDrop complements existing FL methods by adding new client- and server-side components. During client training, convolutional layers are sparsified by interleaving them with channel-wise dropout layers, namely the SyncDrop layers. Section 3.3 discusses in depth how SyncDrop layers compute dropout decisions. Additional optimization stages during server aggregation are added to encourage collaboration between sparse clients. Figure 2 shows how FedDrop extends traditional FedAvg (McMahan et al., 2017). After each round of client training, the server begins by identifying
parameter update directions of each client from the previous round, and computes a cross-client trajectory similarity matrix for each channel neuron. This is then followed by an optimization stage that iteratively minimizes the trajectory disparities by tuning dropout keep probabilities for each channel neuron on each client for the next training round (Section 3.6). This process also takes into consideration the FLOPs budget constraints to sparsify client models (Section 3.4). Finally, following FedAvg, the server broadcasts the weighted-average model parameters to all training clients in the new round, and finally sends the optimized probabilities to the respective clients.

heading: 3.2 PRELIMINARIES AND DEFINITIONS
text: FedDrop complements most FL algorithms. In this paper, we focus on its use on FedAvg (McMahan et al., 2017). We assume the training loss function of a client c ∈ C to be `c(θc), where θc comprises the parameters {θ[l]c : l ∈ L} of all layers in the model of client c. In each FedAvg training round r, clients begin by training on the loss function, with initial parameters θ(r) received from the server for this round:
θ(r+1)c = SGDc ( `c,θ (r), η, E ) . (1)
Here SGDc indicates that client c carries out stochastic gradient descent (SGD) on `c(θ (r) c ) locally, and it uses a learning rate η forE epochs. The FedAvg server then aggregates client model parameters after the rth training round, by taking the weighted average of them:
θ(r+1) = ∑
c∈C λcθ (r+1) c , (2)
where λc is the weight of client c and is proportional to the size of its training set |Dc| with∑ c∈C λc = 1. Finally, the (r + 1) th training round starts by repeating the above procedure.

heading: 3.3 SYNCHRONIZED DROPOUT
text: To induce sparsity in a stochastic manner, and subsequently introduce training correlation across clients, we designed a threshold-based dropout layer, SyncDrop, to synchronize dropout decisions across multiple clients.
Initially for all sampled clients c ∈ C ⊆ C, the server provides them with their corresponding dropout keep probabilities pc ∈ [0, 1]N, and we assume p ∈ [0, 1]C×N to be the concatenated probabilites from all sampled clients. Each element pnc in p denotes the keep probability of the channel neuron n ∈ N in client c ∈ C. We additionally use p[l]c , a slice of pc, to indicate the probabilities of all channels in the lth layer. During training, the lth layer samples t[l]n from a uniform distribution U(0, 1) for all channel neurons n in the layer, where t[l]n is shared across clients by using the same random seed. If pnc < t [l] n , a channel n in client c is dropped, i.e. we set the entire channel map to zero; otherwise, the channel values are scaled by 1/pnc during training 2. Formally, for each client c, the lth dropout layer computes for the input x[l]:
drop ( x[l],p[l]c ) , d[l]c ◦ x[l] ◦ p[l]c ◦−1 ,where d[l]c , 1[t [l] < p[l]c ]. (3)
Here, 1[z] is the (element-wise) indicator function, which is equal to 1 when the condition z is met and 0 otherwise, ◦ refers to the element-wise product, the term p[l]c ◦−1 represents the element-wise inverse of p[l]c , and finally all elements of t
[l] are independently sampled from U(0, 1) and shared across clients. Figure 3a provides a high-level overview of the procedure described above.
From the local perspective of a client c, the dropout decision d[l]c is equivalent to the Bernoulli distributions B(p[l]c ). An important distinction from independent Bernoulli distributions is that the distribution of d[l]c is correlated across clients c ∈ C. Given a pair of clients (i, j) for channel n:
Et∼U(0,1) [ dni d n j ] = ∫ 1 0 1[t < pni ] · 1[t < pnj ]dt = min ( pni ,p n j ) . (4)
This enables us to adjust the correlation of the same channel neurons between any pairs of clients, and Section 3.6 makes use of this property to minimize the model update disparities across all clients.
2This is known as “inverted dropout” and implemented by PyTorch (2021) and TensorFlow (2021). Figure 7 provides ablation results on scaling for dropout to justify the design choice.
We define f̂ to be the sparsified variant of the original model f , where SyncDrop layers are placed immediately after each convolutional layer with ReLU activation (Figure 3b). This enables convolutional layers to be doubly accelerated by taking advantage of sparsities at both ends and skipping dropped channels in input/output feature maps. Recalling the client model loss function `c, we additionally use ˆ̀c to denote the accelerated variant of the sparse model f̂ with probabilities pc.
It is noteworthy that active channels are sampled for each local training mini-batch. In addition, the synchronization of dropout decisions is carried out with identical random seeds for all clients, and thus it eliminates the need to communicate between clients. Finally, for inference, one may choose to enable the SyncDrop layers for speed, or skip these layers entirely. In our evaluations, SyncDrop layers are disabled for improved test accuracies.

heading: 3.4 FLOPS CONSTRAINTS
text: To encourage sparsity in models and for the optimization of p, we adopt a hyper-parameter r ∈ (0, 1] that adjusts the FLOPs budget ratio, and define a shared global FLOPs-budget constraint g(p) ≥ 0 for all participating clients C ⊆ C in one round, where:
g(p) = r − ∑
c∈C flops(f̂ ,pc)
/ (|C|flops(f)), (5)
and the terms flops(f̂ ,pc) and flops(f) respectively denote the FLOPs of a sparse model with probabilities pc, and the FLOPs of the dense model f . Moreover, this constraint can be easily modified to allow client- and layer-wise, and even heterogeneous budgets by summing over targets with finer granularity. In Appendix B, we explain how the FLOPs of a model is computed.

heading: 3.5 PARAMETER INITIALIZATION
text: At model initialization, we set all values in p to be a uniform constant p such that g(p) = 0. Since the training clients are sparse at initialization, the variance of each initialized parameters θ[l]ij ∈ θ[l] of the lth layer must satisfy the following condition: Theorem 1 (Parameter initialization). To avoid vanishing/exploding gradients at initialization, the variance of the parameters of the lth layer followed by a dropout layer with keep probability p should be 2p/N , where N = C [l−1]K [l−1]2 is the size of the receptive field of each output channel.
In this paper, we introduce a modified He initialization (He et al., 2015), that is θ ∼ N (0, √
2p/N). In Appendix C.1 we provide a proof of this theorem. Both Hendrycks & Gimpel (2016) and Pretorius et al. (2018) derived similar schemes for fixing parameters initialization of layers followed by dropout. Figure 7 further justifies this decision empirically by ablation of the fixed initialization.

heading: 3.6 OPTIMIZING DROPOUT KEEP PROBABILITIES
text: In our experiments, we found that using a constant dropout keep probability is often detrimental to the joint optimization of computation and communication costs. To speed up training while minimizing the impact of model sparsity on global training convergence, we therefore aim to optimize the dropout keep probabilities p for the following objective:
minp E ∥∥∥∑c∈Cλc∇θc ˆ̀c −∑c∈Cλc∇θc`c∥∥∥2
2 , (6)
i.e. the mean square error between the averaged gradients of all training clients with and without SyncDrop layers. Intuitively, it aligns the parameter updates of sparse models with the original dense variants for an improved convergence behavior. As a result of the correlated dropouts across clients (4), the objective above can be minimized by adjusting all dropout keep probabilities pnc across neuron n on client c.
As it is impractical to communicate each step of SGD across clients, FedDrop instead optimizes an alternative objective derived in the theorem below (proved in Appendix C.2) that approximately minimizes the original: Theorem 2 (Approximate Objective). Assuming that for a channel neuron n ∈ N and a pair of clients i, j ∈ C, we have Ŝnij , λiλj max(pni ,pnj )∆θni · ∆θnj , where ∆θni ,∆θnj are the parameter updates after a round of FL training, pni ,p n j are the current dropout keep probabilities, and max(x,y) represents the element-wise max between x and y, the optimization objective of (6) can be approximated by:
minq obj(Ŝ,q) , minq ∑ n∈N,i,j∈C Ŝ n ij/max(q n i ,q n j ). (7)
The optimized values q can subsequently be used as the new dropout keep probabilities for the next round of FL training.
It is also desirable to have a bounded optimization to avoid extreme solutions. Without a FLOPs constraint, both objectives (6) and (7) can be trivially minimized with p = 1, i.e. when none of the channels are dropped, and Appendix C.2 provides the proofs of the following theorem: Theorem 3 (Optimization bound). The objective minq ∑ nij Ŝ n ij/max(q n i ,q n j ) is trivially minimized
to ∑
nij Ŝ n ij when q n c = 1 for ∀n ∈ N, c ∈ C.
To enforce sparsity, we therefore further constrain the optimization of q to be within the feasible set g(p) ≥ 0 as defined in Section 3.4. In practice on the server-side after gathering model updates from a round of sparse client training, we minimize the overall objective via gradient descent with the interior point method: minq obj(Ŝ,q)− µ log(g(q)). (8) The term−µ log(g(q)) constrains the solution q to be within the feasible set g(q) ≥ 0, and the hyperparameter µ tunes the strength of the regularization, and is kept constant at 10−4. In Appendix D, we describe the overall FedDrop algorithm in depth, and discuss the computational and memory implications of FedDrop.

heading: 4 EXPERIMENTAL RESULTS
text: In this section, we present a comprehensive evaluation of both the communication and computation costs of FedDrop and compare it against other FL algorithms, namely FedAvg (McMahan et al., 2017), FedProx (Li et al., 2020), SCAFFOLD (Karimireddy et al., 2020), and Caldas et al. (2018a). We also introduce UniDrop for reference — a simple channel-wise dropout with a uniform keep probability that is kept constant as the comparison baseline.
We conduct experiments on three popular open datasets CIFAR-10 (Krizhevsky et al., 2014), FashionMNIST (Xiao et al., 2017) and SVHN (Netzer et al., 2011). To simulate non-IID training data distribution with a concept shift, for every class label we split the training dataset into |C| parts to be respectively received by the clients, where the size of all parts follow the Dirichlet distribution D|C|(α) (Hsu et al., 2019; Wang et al., 2020b). For the full participation case, we used α = 0.5, 20 clients for CIFAR-10, and 100 clients for Fashion-MNIST and SVHN in all experiments by default.
We used a 9-layer VGG-style architecture (Simonyan & Zisserman, 2015) for CIFAR-10, wheras Fashion-MNIST used a LeNet (Lecun et al., 1998) model and SVHN a smaller variant, and we provide the details of the models and datasets in Appendix E.
For a fair comparison, we performed a grid-based search on the remaining hyperparameters for baseline algorithms to find configurations that provide the best accuracy after 100 rounds of training for each dataset. The search results suggested that typically setting the batch size B = 4 and learning rate η = 0.02 were universally optimal for most methods under comparison. FedProx additionally introduced a regularizer µ which was also optimized within {0.0001, 0.001, 0.01, 0.1}. Finally, FedDrop introduces a global FLOPs budget per round as a ratio between r ∈ (0, 1] that can be adjusted to steer the communication/computation trade-off as defined in (5), for which we specify the values in the experiments below.
FLOPs vs. Accuracy. We first focus on the amount of FLOPs required by the algorithms, and provide a set of comparisons across the FL methods for the same number of local epochs per round E, given their respective optimal hyperparameters with an unbounded communication budget (Figure 4). Here, we used a FLOPs ratio r = 0.5 for FedDrop. Frequent communication rounds allow the server to more frequently optimize dropout keep probabilities for FedDrop, which can notably improve convergence under the same FLOPs budget. In contrast, other methods struggle to improve their convergence speed.
FLOPs vs. Communications. To bolster the trade-off between communication and computation, we additionally conducted a hyperparameter exploration to navigate the Pareto frontiers of this trade-off relationship. The hyperparameters include the number of local epochs per round E and the batch size B. In addition, FedDrop can search for the Pareto optimal r in the set {0.2, 0.25, 0.333, 0.5}. Figure 5 shows the FLOPs/communication trade-off of the FL methods. It is notable that with decreasing FLOPs budgets, FedDrop expends minimal communications, whereas other methods may have significantly larger communication costs to reach the same target accuracy. Conversely, in Table 1, we show that FedDrop can be up to 3× more efficient in terms of FLOPs under the same communication budget for the same accuracy goal.
Fractional Device Participation. To measure the performance of FedDrop in more practical distributed scenario, we scale the number of devices to 1000, while reducing the fraction of participation rate φ per round to only 1%, yielding results in Table 2. We refer readers to Appendix F.5 for more detailed trade-off results.
Dataset Distribution Imbalance. To investigate how different data distributions affect the performance of FedDrop rigorously, we provide a sweep of α ∈ {5, 0.5, 0.05} for the CIFAR-10 dataset splits D|C|(α) to simulate increasing degrees of distribution imbalance in Figure 6. It shows that FedDrop is the best method at maintaining accuracy with distribution imbalance w.r.t. the FLOPs budget. We believe SCAFFOLD may be unsuited for larger models as it diverged under these scenarios.
Ablation Study. To verify our design choice, we conduct a thorough set of ablations in Figure 7. We observe a certain degree of accuracy drop as these components are ablated, confirming the significance
of synchronization. This then further justifies the effectiveness of dropout scaling and initialization that comes with theoretical guarantees.
20 40 60 80 100 60
65
70
75
80
85
Number of Communication Rounds
A cc
ur ac
y (%
)
50 k 10 k 5 k 1 k
(a) Optimization iterations.
20 40 60 80 100 60
65
70
75
80
85
Number of Communication Rounds
A cc
ur ac
y (%
)
0.01 0.001 0.0001 0.00001
(b) FLOPs regularizer.
1014 1014.2 1014.4 1014.6 1014.8 50
60
70
80
Number of FLOPs
A cc
ur ac
y (% ) 0.75
0.5 0.333 0.25
(c) FLOPs vs. accuracy.
0 20 40 60 80
20
40
60
80
Number of Communication Rounds
A cc
ur ac
y (%
)
0.75 0.5 0.333 0.25
(d) Rounds vs. accuracy.
Figure 8: Sensitivity analyses of (a) optimization iterations I , (b) FLOPs constraint regularizer µ, and (c, d) the FLOPs budget ratio r.
Sensitivity Analysis. Figure 8 produces a sensitivity analysis of the hyperparameters introduced by our server-side probability objective minimization (7), namely the number of gradient descent iterations I and the strength of regularization µ. It turns out that additional iterations can improve convergence, which is desirable as the server compute budget is usually not the performance bottleneck. Moreover, a broad range of regularization µ do not have an adverse impact on accuracy, eliminating the need for a search. We also present an analysis on the effect of the FLOPs ratio r as defined in (5). Reducing this ratio can sparsify models, for a faster initial convergence w.r.t. FLOPs. However, it comes at a cost of a slower convergence in latter rounds.
Additional Results. In Appendix F, we provide additional results and analyses. Specifically, Appendix F.1 includes the trade-off curves for all datasets under different accuracy budgets. We examine the performance impact on FL algorithms with different degrees of data imbalance (Appendix F.2), and varying hyperparameters such as local epochs per round E (Appendix F.3) and the FLOPs budget ratio r (Appendix F.4). Finally, Appendix F.5 provides full results with fractional device participation.

heading: 5 CONCLUSION
text: Recent federate learning (FL) algorithms focus on the reduction of communication costs, while neglecting the expensive local compute required by edge clients. This could be further exacerbated as the models we employ may increase in size over time to attain a higher task performance. We presented FedDrop, a novel FL technique that uses synchronous dropout with adaptive keep probabilities to concentrate the clients’ training effort on neurons that they specialize well, while making the models sparse to reduce computational costs. Our experiments show that FedDrop can push the Pareto frontiers of communication/computation trade-off of FL scenarios notably further than existing algorithms. We believe this paves the way for future work that allow FL algorithms to scale up models being trained considerably, and consequently improve their performance on more challenging training tasks.

heading: 6 ETHICS STATEMENT
text: We believe FedDrop can have a positive impact on the environment as it aims to make federated learning more efficient.
The federated learning algorithms require communicating locally optimized client models to global server. In such a scenario, training data could potentially be recovered from model updates (Yin et al., 2021; Wang et al., 2020a), and this could have privacy implications. Although our work do not empirically address the privacy concern, we believe the added dropout layers do not bring noticeable impact to the underlying FL algorithm, for the reasons below:
• The dropout probabilities are computed on the server with client updates, the optimized probabilities are simply derived from information that are already collected by the base algorithm.
• There is no inter-client information exchange or leakage, as each client only receives its corresponding dropout probabilities.

heading: 7 REPRODUCIBILITY STATEMENT
text: We provide detailed information about model configuration in Appendix E, dataset splitting and partitioning in Section 4, and hyperparameters for optimization in Section 4. We try to justify the design choices of our method with ablation in Section 4, and all the experiments reported in this manuscript are conducted exhaustively with broad choices of hyperparameters to ensure fair evaluation for all methods. The code would be available during the review period and made public upon acceptance.

heading: A TABLE OF NOTATIONS
text: For your convenience, we provide a list of notations used in this paper in Table 3.

heading: B COMPUTING THE NUMBER OF FLOPS
text: In a sparse layer f̂ [l], the expected number of FLOPs per image-step is the sum of the flops required by both the convolution and ReLU activation:
flops(f̂ [l]) = 2(Ĉ [l]Ĉ [l−1]K [l] 2 H [l]W [l] + Ĉ [l]H [l]W [l]), (9)
where K2 is the 2-dimensional kernel size, and H ×W is the output feature map size. Moreover, Ĉ [l−1] and Ĉ [l] are the number of active input and output channels respectively; we expect Ĉ [l], the average number of remaining output channels for layer l, to be C [l] mean ( p[l]c ) , where mean(z) takes the mean of elements in z and C [l] is the number of total output channels of layer l.
To generalize, we can rewrite the total number of FLOPs required by the overall model f̂ to be: flops(f̂ ,pc) = 2 ∑L l=0C [l] mean ( p[l]c )( K [l] 2 C [l−1] mean ( p[l−1]c ) + 1 ) H [l]W [l], (10)
and assume p[0] = 1 and p[L] = 1, since both the input and output of the model are considered to be dense. Finally, by setting p = 1 for the entire model, we can get the number of FLOPs consumed by the overall model per image.

heading: C PROOFS

heading: C.1 MODEL PARAMETER INITIALIZATION
text: Here, we restate Theorem 1 from Section 3.5: Theorem 1 (Parameter initialization). To avoid vanishing/exploding gradients at initialization, the variance of the parameters of the lth layer followed by a dropout layer with keep probability p should be 2p/N , where N = C [l−1]K [l−1]2 is the size of the receptive field of each output channel.
The theorem proof depends on the following definitions and assumptions. Definition 1. A convolutional layer with channel dropout decisions d[l] can be written as:
x[l+1] , relu(z[l]), where z[l] , p[l]c ◦−1 ◦ d ◦ (ϑ[l]x[l] + b[l]), (11)
where we flatten the input features and the convolutional weights to be respectively z[l],x[l] ∈ RN×1 and ϑ[l] ∈ RC[l+1]×N [l] , with N [l] = K [l]2C [l] being the size of the input receptive field. Recall that ◦ indicates the element-wise product, and the notation z◦−1 is the element-wise inverse of z. Here each dropout decision in d[l] is sampled with:
d[l] , 1[t[l] < p[l]]. (12)
At initialization, we assume that p[l] shares a constant dropout keep probability p, and we can immediately derive that for any channel neuron i at initialization:
E[d[l]i ] = p · 1 + (1− p) · 0 = p, (13) var(d[l]i ) = E[d [l] i 2 ]− E[d[l]i ]2 = p · 1 + (1− p) · 0− p2 = p− p2. (14)
Assumption 1. We assume the model input x[0] and for each layer l the intermediate values z[l] to be IID and all have zero-valued means. We also let the individual weight parameters ϑ[l] to be initialized with IID and zero means, and the biases b[l] = 0. Therefore, for each layer l ∈ [1, L]:
E(z[l]ij) = 0 for z [l] ij ∈ z[l], (15) E(ϑ[l]ij) = 0 for ϑ [l] ij ∈ ϑ[l]. (16)
Assumption 2. To avoid vanishing/exploding gradients, we restrict the variances of each value in z[l] to be constant across all layers l ∈ [1, L], i.e.
var(z[m]i ) = var(z [n] j ) for m,n ∈ [1, L], z[m]i ∈ z[m], z[n]j ∈ z[n]. (17)
We adapt the proofs from (He et al., 2015; Glorot & Bengio, 2010) by taking into consideration the variance introduced by dropout layers:
Proof. By Definition 1, for any value z[l]i ∈ z[l], at initialization: var(z[l]i ) = 1 p2 var( ∑ j d [l] i ϑ [l] ijx [l] j + d [l] i b [l] i )
As d[l]i , ϑ [l] ij and x [l] j are independent, and b [l] = 0 at initialization following Assumption 1:
= N [l]
p2 var(d [l] i ϑ [l] ijx [l] j )
= N [l]
p2 ( var(d[l]i ) + E[d [l] i ] 2 )( var(ϑ[l]ij) + E[ϑ [l] ij ] 2 )( var(x[l]j ) + E[x [l] j ] 2 )
− E[d[l]i ]2 E[ϑ[l]ij ]2 E[x[l]j ]2. (18) Again following Assumption 1, E[ϑ[l]ij ] = 0. Also, var(x [l] j ) + E[x [l] j ] 2 = E ( x[l]j 2), E[d[l]i ] = p, and var(d[l]i ) = p(1− p):
var(z[l]i ) = N [l]
p var(ϑ[l]ij)E
( x[l]j 2) . (19)
As proven by (He et al., 2015) for ReLU-nonlinearity, E ( x[l]j 2) = 12 var(z [l−1] j ), and using Assumption 2, we can let var(z[l]i ) = var(z [l−1] j ), and solve (19) to give var(ϑ [l] ij) = 2p N [l]
. Finally, if we are using Gaussian distributions to initialize model parameters, then ϑ[l]ij ∼ N ( 0, √ 2p/N [l] ) .

heading: C.2 OPTIMIZATION OBJECTIVE
text: This section provides proof sketches for the optimization objective (7) described in Section 3.6. Let’s assume gc , λc∇θ`c, and ĝc , λc∇θ ˆ̀c to simplify the derivations below. Lemma 1. The sample gradient with dropout ĝc can be approximated by dc ◦ p◦−1c ◦ gc.
Proof. Consider the 1st-order Taylor expansion of ˆ̀c(θ): ˆ̀ c(θc) = `c(dc ◦ p◦−1c ◦ θc) ≈ `c(θc) + ( ∇θc`c )ᵀ( dc ◦ p◦−1c ◦ θc − θc ) .
Differentiating both sides by θc, we have the following, where H ( `c ) signify the Hessian of `c(θc):
∇θc ˆ̀c −∇θc`c ≈ ∇θc (( ∇θc`c )ᵀ( dc ◦ p◦−1c − 1 ) ◦ θc ) = ( ∇θc`c ) ◦ ( dc ◦ p◦−1c − 1 ) + ( (dc ◦ p◦−1c − 1) ◦ θc )ᵀ H(`c).
We omit the second term ( (dc ◦ p◦−1c − 1) ◦ θc )ᵀ H(`c) for simplicity, as the Hessian of `c(θc) is compute-intensive, thus: ĝc − gc = λc ( ∇θc ˆ̀c −∇θc`c ) ≈ gc ◦ ( dc ◦ p◦−1c − 1 ) . (20)
In practice, we cannot obtain Snij , E [ gni g n j ] , as evaluating the gradients of dense models defeats the intention of accelerated training. We can instead estimate the value with the following lemma: Lemma 2. E [ gni g n j ] can be approximated with E [ ĝni ĝ n j ] max(pni ,p n i ), where p n i ,p n i are the dropout keep probabilities from the previous round of the neuron n of respective clients i, j.
Proof. From Lemma 1, by rearranging (20), we have gc ≈ ĝc ◦ d◦−1c ◦ pc for client c. Therefore, E [ ĝni ĝ n j ] ≈ E [gni gnj dni dnj pni p n j ] As dc and gc are independent for each client c,
= E [ gni g n j ] E [dni dnj pni p n j ] By the definition of the SyncDrop layer in Section 3.3,
= E [ gni g n j ] ∫ 1 0 1[t≤pni ] 1[t≤p n j ]
pni p n j
dt = E [ gni g n j ]min(pni ,pnj ) pni p n j .
Rearrange to give: E [ gni g n j ] = E [ ĝni ĝ n j ] max(pni ,p n j ). (21)
Finally, we reiterate Theorem 2 that introduces the approximate objective (7):
Theorem 2 (Approximate Objective). Assuming that for a channel neuron n ∈ N and a pair of clients i, j ∈ C, we have Ŝnij , λiλj max(pni ,pnj )∆θni · ∆θnj , where ∆θni ,∆θnj are the parameter updates after a round of FL training, pni ,p n j are the current dropout keep probabilities, and max(x,y) represents the element-wise max between x and y, the optimization objective of (6) can be approximated by:
minq obj(Ŝ,q) , minq ∑ n∈N,i,j∈C Ŝ n ij/max(q n i ,q n j ). (7)
Proof. We start by deriving from (6): minp E ∥∥∥∑c∈Cλc∇θc ˆ̀c −∑c∈Cλc∇θc`c∥∥∥2
2 = E [(∑ c∈C(ĝc − gc) )ᵀ(∑ c∈C(ĝc − gc) )]
= E [∑
i,j∈C(ĝi − gi) ᵀ (ĝj − gj)
] ,
following Lemma 1, ≈∑i,j∈C E[(gi ◦ (di ◦ q◦−1i − 1))ᵀ(gj ◦ (dj ◦ q◦−1j − 1))] = ∑ i,j∈C,n∈N E [ gni g n j (dni qni − 1 )(dnj qnj − 1 )] ,
as dc and gc are independent for each client c, and following Lemma 2,
= ∑ i,j∈C,n∈N E [ ĝni ĝ n j ] max(pni ,p n j )E [(dni qni − 1 )(dnj qnj − 1 )] ,
by definition of the SyncDrop layer in Section 3.3,
= ∑ i,j∈C,n∈N E [ ĝni ĝ n j ] max(pni ,p n j ) ∫ 1 0 ( 1[t≤qni ] qni − 1 )( 1[t≤qnj ] qnj − 1 ) dt
= ∑ i,j∈C,n∈N E [ ĝni ĝ n j ] max(pni ,p n j ) (
1 max(qni ,q n j ) − 1 )
(22)
Finally, as it is impractical to use single step gradients ĝni and ĝ n j on the server, we approximate them with the respective parameter updates after one round of training, i.e. λi∆θni and λj∆θ n j ; recall the definition of Ŝ, we have
≈∑i,j∈C,n∈N Ŝnijmax(qni ,qnj ) −∑i,j∈C,n∈N Snij . (23) The first term is exactly (7), and the second term is a constant and can be safely removed from the optimization objective.
The value Ŝnij can be considered as the similarity of the model updates between clients i and j for the group neuron n. Intuitively, from the perspective of an individual neuron n and the two clients i and j, minimizing the term Ŝnij/max(p n i ,p n j ) incentivizes p n i and p n j to grow when Ŝ n ij is positive, and to become smaller when it is negative.
With the derivations above, we can proceed to prove the objective in Theorem 2 is bounded with a minimum. First, we can show that the matrix Sn is positive semidefinite for each neuron n ∈ N. Lemma 3. The matrix Sn is positive semidefinite for all n ∈ N.
Proof. By definition, Sn = E[gngnᵀ], where gn comprises the gradient vectors of the nth neuron across clients c ∈ C. The matrix Sn is therefore positive semidefinite, as it is an averaged sum of samples of positive semidefinite matrices gngnᵀ.
Theorem 3 (Optimization bound). The objective minq ∑ nij Ŝ n ij/max(q n i ,q n j ) is trivially minimized
to ∑
nij Ŝ n ij when q n c = 1 for ∀n ∈ N, c ∈ C.
Proof. Without loss of generality, we can permute the matrix Ŝn and the probability vector pn to assume a descending order on the probabilities. Let Tn , ΠŜnΠᵀ and πn , Πpn denote the respective permuted trajectory similarity and probability variants, where Π ∈ {0, 1}|C|×|C| is the permutation matrix, such that∑
nij Ŝnij max(pni ,p n j )
= ∑
nij Tnij max(πni ,π n j )
and πn1 ≥ πn2 ≥ · · · ≥ πn|C|. (24) We introduce a boolean vector hk ∈ R|C|, where each element hkc of it is defined as:
hkc , 1[π n k ≥ πnc ],
we then let Hk , hkhkᵀ, and it follows that Hkij = h k i h k j = 1[π n k ≥ πni ] 1[πnk ≥ πnj ]. Hence, we can rewrite:∑ nij Tnij max(πni ,π n j ) = ∑ nij T n ij min ( 1 πni , 1πnj ) = ∑
nij T n ij (∑|C|−1 k=1 Hkij−H k+1 ij
πnk +
H |C| ij πn|C| ) = ∑
nij T n ij (∑|C|−1 k=1 Hkij πnk −∑|C|−1k=1 Hk+1ijπnk + H|C|ijπn|C| ) = ∑
nij T n ij ( H1ij πn1 + ∑|C| k=2 H k ij ( 1 πnk − 1πnk−1 )) = ∑
n (∑ ij Tnij πn1 + ∑|C| k=2 ∑ ij T n ijh k i h k j ( 1 πnk − 1πnk−1 )) = ∑
n ( 1ᵀTn1 πn1 + ∑|C| k=2 h kᵀTnhk ( 1 πnk − 1πnk−1 )) . (25)
It is notable that 1πnc ≥ 1 for all c ∈ C. The first term is bounded by a non-negative 1 ᵀTn1 i.e. the sum of all elements in a positive semidefinite matrix. The second term is also non-negative, as hk
ᵀ Tnhk ≥ 0 by the definition of a positive semidefinite matrix Tn, and 1πnk − 1 πnk−1 ≥ 0 by the ordering proposed in (24). Notably if πnc = 1 for all n ∈ N and c ∈ C, then the objective is trivially minimized, with a minimum value ∑ nij T n ij = ∑ nij S n ij . We therefore include the FLOPs budgets discussed in Section 3.4 and Appendix B into our optimizations to solve for non-trivial p.

heading: D THE FEDDROP ALGORITHM

heading: D.1 HIGH-LEVEL OVERVIEW
text: Algorithm 1 provides an overview of the FedDrop FL algorithm. It accepts the client loss functions ˆ̀ c with SyncDrop layers, client data weights λc, the FLOPS constraints function g : [0, 1]
C×N → R, the SGD learning rate η, the number of local epochs E the number of FL rounds R, and a client sub-sampling ratio φ. It returns optimized model parameters θ(R+1) on the final round.
The algorithm starts by initializing model parameters θ(0) to be shared across all clients, and assigns a uniform keep probability to p for all SyncDrop layers, which satisfies the FLOPs constraint g(p) = 0 (lines 2 and 3). After collecting trained model parameters θ(r+1)c from the current round r + 1 (line 6), the FedDrop server computes ∆θ(r+1)c , i.e. the difference between the averaged model from the previous round θ(r) with the trained model θ(r+1)c for client c (line 7, also in Figure 2a).
It then measures the trajectory similarity Ŝnij (r+1) for each neuron n between each pair of clients (i, j), as described on line 11 and illustrated in Figure 2b. The term λi∆θni (r+1) · λj∆θnj (r+1) computes
the dot-product between the weighted parameter updates of clients i and j, and max ( pni (r),pnj (r) ) adjusts for the covariance from the synchronized dropouts of the previous round. Finally, the objective argminq obj(Ŝ,q) is optimized to minimize the impact of model sparsity on convergence.

heading: D.2 COMPUTATIONAL OVERHEAD OF THE OPTIMIZATION OBJECTIVE
text: The complexity of computing the trajectory similarity in each round is O(W |C|2), where W is the number of parameters and |C| is the number of participating clients. Additionally, the probability
Algorithm 1 The FedDrop algorithm.
1: function FEDAVGDROP({(ˆ̀c, λc) : c ∈ C}, g, η, E,R, φ) 2: random_initialize(θ(0),p) . Initialize parameters and a uniform keep probability. 3: for r ← 1, 2, . . . R do 4: C← subsample(C, φ) . Sample a subset of clients. 5: par for c ∈ C do 6: θ(r+1)c ← SGDc ( ˆ̀ c,θ (r),pc, η, E )
. Client training with SyncDrop layers. 7: ∆θ(r+1)c ← θ(r) − θ(r+1)c . Client update trajectories. 8: end par for
9: θ(r+1) ←∑c∈C λcθ(r+1)c . Server aggregation with FedAvg. 10: for i ∈ C, j ∈ C, n ∈ N do . Trajectory-similarity estimation. 11: Ŝnij ← λiλj max ( pni ,p n j ) ∆θui
(r+1) ·∆θuj (r+1) 12: end for . Optimal dropout probabilities. 13: p← argminq obj(Ŝ,q) s.t. g(q) ≥ 0 . Objective from (7) 14: end for 15: return θ(R+1) 16: end function
optimization stage on the server has a complexity of O(I|N||C|2), where I is the number of optimization iterations, |N| is the number of channels in the model. As an example, the VGG-9 model with 20 clients and I = 10k steps of optimization would require only 53 G FLOPs for computing the trajectory similarity, and ≤ 448 G FLOPs for the optimization steps, which can often exit early because it has converged. In practice, the server optimization only takes less than 1 minute on average. For reference, if we use local epochs per round E = 8, and a FLOPs ratio r = 0.5, the 20 clients require a total of 25.8 T FLOPs per round. Overall, under this setting the additional computation overhead incurred by the server is only about 2% of total computational costs. For more practical consideration, it is noteworthy that typically, the server is much less constrained in terms of computational power compared to the edge clients. In addition, given the above example, a single client requires 1.29 T FLOPs and the server requires 448 G FLOPs in total. The former consumes multiple hours of training in a realistic implementation on an iPhone 12 Pro (extrapolating from the profiling result in Table 10), whereas the latter only requires minutes to complete on an Nvidia V100 GPU.

heading: D.3 AN EXAMPLE HISTORY OF PROBABILITY DISTRIBUTION
text: To better understand the probability optimization, we plot the histogram of optimized probabilities for each training round in Figure 9. It reveals that after multiple rounds of training, FedDrop gradually learns the importance of channels and start to make certain channels contribute less to the model training. Note that the histograms only provides distributions of probability values of all clients, each client may have vastly different optimized probabilities for the same channels.

heading: D.4 MEMORY REDUCTION
text: As SyncDrop layers can sparsify models, many channels in the activation map can simply be skipped and never stored in both the forward and backward passes. For this reason, FedDrop can further reduce the memory cost associated with training. As an example, VGG-9 equipped with SyncDrop can see up to 40% reduction in the memory requirement for training. We summarized the memory reduction for different scenarios in Table 4.

heading: D.5 ACCURACY OF FLOPS ESTIMATION
text: To confirm the accuracy of our FLOPs estimation with the expected values, in Table 5 we sample the dropout layers in the Fashion-MNIST model with FLOPs budget ratios r ∈ {0.25, 0.5, 0.75}, can compute the actual FLOPs counts. Under the setting of R = 100 training rounds and local epochs

heading: 4 3.02 5.82 8.83 0.75 8.08 8.530.5 7.33 17.07

heading: 8 6.03 5.82 11.85 0.75 10.34 12.720.5 8.83 25.45
text: per round E = 8. We found the error to be within 0.06%. The low deviation of the expected value from the actual result is an effect of the law of large numbers.

heading: E EXPERIMENT CONFIGURATIONS
text: Models. Here, we provide the layout of the models used by the respective datasets. All models listed for computer vision (CV) tasks (in Tables 6 to 8) are feed-forward CNNs. Table 9 shows the character-level single-layer LSTM model used for the natural language processing (NLP) task. Note that sparsity for the LSTM model were introduced for the hidden state and memory cell neurons.
Datasets. The experiments are carried out on three popular CV datasets (CIFAR-10 (Krizhevsky et al., 2014), Fashion-MNIST (Xiao et al., 2017) and SVHN (Netzer et al., 2011)), and a natural language dataset (Shakespeare (Caldas et al., 2018b)). We use the standard train/test split on these datasets and report accuracy on test dataset with a global model on server-side. To simulate different heterogeneity levels, for CV tasks we split the datasets by class label with the distribution D|C|(α) (Hsu et al., 2019; Wang et al., 2020b), and adjust the α constant for varying the degree of data heterogeneity, α = 0.5 unless specified for all experiments. Simple data augmentation was also employed for the training images, including random crop and normalization, and random flip in addition for CIFAR-10. For the Shakespeare NLP dataset, we follow LEAF (Caldas et al., 2018b) for next-character prediction, and consider each speaking role in each play as an individual client, and the resulting splits are thus inherently heterogeneous since each local client only contains the corresponding role’s sentences.
Hyperparameters. We consider different hyperparameter settings for different tasks:
• Client Sampling. For full client participation, we use 20 clients for CIFAR-10 and 100 clients for the remaining CV and NLP tasks. To simulate the fractional client participation, we randomly sample 10 out of 1000 clients on all datasets for CV tasks and 10 out of 660 clients on Shakespeare dataset for NLP task.
• Optimization. As for CV tasks, we searched the hyperparameters and found setting batch size B = 4 and learning rate η = 0.02 were universally optimal for most methods. For NLP task on Shakespeare dataset, the batch size B = 32 and learning rate η = 0.2 were similarly found and made constant for all methods under comparison.
We noticed that the reported accuracy could be further improved if deeper models are used for baselines as they can exhibit greater redundancy and thus benefit the performance of FedDrop. However, we are not pursuing state-of-the-art accuracy on these datasets but instead comparing FedDrop with other methods under fair experimental settings and validating its effectiveness. We also noticed that there exist larger alternative datasets (Hsu et al., 2020) for evaluation. Nevertheless, evaluating and reproducing all competing FL methods on such a large dataset is infeasible given the limited (university) computational budgets. Alternatively, we aim to conduct a fair empirical evaluation that demonstrate the effectiveness of FedDrop under a broad variety of scenarios.

heading: F ADDITIONAL EXPERIMENTAL DETAILS AND RESULTS

heading: F.1 TRADE-OFF CURVES
text: Figures 10 to 12 find the optimal configurations for each competing FL methods that produces their corresponding communication/computation trade-off curves. It is notable that FedDrop can extend the trade-off relationships much further and shows Pareto dominance against other algorithms under most conditions. In these figures, the configurations explored include combinations of FLOPs budget ratios r ∈ {0.25, 0.333, 0.5, 0.75} and local epochs per round E ∈ {1, 2, 4, 8, 16}.

heading: F.2 DATA IMBALANCE
text: Figures 13 to 15 examine the effect of data imbalance on the performance of FedDrop. Again, it shows that FedDrop can manage the fastest convergence when compared to other methods, and an increase in data imbalance further enhances its effectiveness against the others.
F.3 VARYING LOCAL EPOCHS PER ROUND
Figures 16 and 17 adjust the number of local epochs per round, and observe that with more frequent rounds (lower E), the accuracy gaps between FedDrop and the other methods become greater.
F.4 VARYING FLOPS BUDGET RATIO
Figures 18 to 20 explore varying FLOPs budget ratios r ∈ {0.75, 0.5, 0.333, 0.25}. We found that a higher sparsity (lower ratio) can improve the overall convergence rate w.r.t. FLOPs expended. However, the final converged accuracy may degrade; this can be fixed by simply adjusting the budget ratio to approach 1.

heading: F.5 FRACTIONAL DEVICE PARTICIPATION
text: We increase the number of all devices to 1000, and reduced the fraction of participating devices per round to 1% for each of the datasets. All competing methods perform uniform device subsampling each round as this was commonly reported for their original experiments McMahan et al. (2017); Li et al. (2020); Karimireddy et al. (2020); Caldas et al. (2018a). For FedDrop to perform well, we uniformly sample devices every 2 rounds, to train with freshly updated probabilities. In Figure 21, Figure 22, Figure 23 and Figure 24, we compare FedDrop against all competing methods, and show that it can bring notable improvements in both communication and computation efficiencies.

heading: F.6 ACTUAL TRAINING RUN TIME
text: In Table 10, we provide the average training times per image on an iPhone 12 Pro running the FashionMNIST model with a batch size of 4. Here, each result was obtained under optimal conditions where no frequency throttling occurred. The structure of the model can be found in Table 6 of Appendix E. Note that we were only able to report CPU run times, as Core ML (cor, 2021) currently only support training on CPUs. We implemented a custom convolutional layer to take advantage of the sparse input and output feature maps.
For the FedAvg motivation example in Section 1, we assumed 20 clients, with 2500 examples per epoch on each client, and 4 training epochs per round, which is in line with our experiments with 20 devices. We extrapolate the results to 100 epochs for a batch size of 4 to estimate the run time to be 983.45µs× 2500× 4× 100 = 16.39 min without frequency throttling. As each parameter is stored as a 32-bit float, each client would upload and download a model of size 226 k × 4 B = 904 kB once per round; the overall transmission per client is thus evaluated to be 904 kB× 2× 100 = 176.56 MB.

heading: F.7 ADDITIONAL RESULTS ON THE MOTIVATING EXAMPLE
text: For reference, in Figure 25 we included all parameter update magnitudes of the first convolutional layer for the motivating example given in Figure 1 in Section 1.

heading: F.8 RESULTS ON THE NLP TASK
text: We provide results on Shakespeare dataset for full client participation in Figure 26 to Figure 28 and fractional client participation in Figure 29 to Figure 31. The experiment configurations are introduced in detail in Appendix E.

Review summaries:
The paper proposes a new method where local workers will drop part of their model using a shared dropout probability received from the server at each communication round. The dropout probabilities are computed by solving an optimization to promote similarity across agent's update. Compared with popular baselines in federated learning, the proposed method demonstrates better performance in terms of number of FLOPs.

The authors propose a new method of coordinated, per client and weight dropout for federated learning. The intuition behind the method is to increase the dropout probability (probability of having a weight set to zero) for weights where different clients often have opposite parity gradients (as when those gradients are summed together are likely to cancel each other out). The authors show that doing this can reduce the computational cost on clients of FL by up to 3x.