SpaFL: Communication-Efficient Federated Learning With Sparse Models And Low Computational Overhead

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:08:55

Abstract: The large communication and computation overhead of federated learning (FL) is one of the main challenges facing its practical deployment over resource-constrained clients and systems. In this work, SpaFL: a communication-efficient FL framework is proposed to optimize sparse model structures  with low computational overhead. In SpaFL, a trainable threshold is defined for each filter/neuron to prune its all connected 
parameters, thereby leading to structured sparsity. To optimize the pruning process itself,  only thresholds are communicated between a server and clients instead of parameters, thereby learning how to prune. Further, global thresholds are used to update model parameters by extracting aggregated parameter importance. The generalization bound of SpaFL is also derived, thereby proving key insights on the relation between sparsity and performance. Experimental results show that SpaFL improves accuracy while requiring much less communication and computing resources compared to sparse baselines. The code is available at https://github.com/news-vt/SpaFL_NeruIPS_2024

Main paper: 
 background and state-of-art: 
 2.1 PERSONALIZED FL Personalized FL is a field of FL whose goal is to produce personalized models for each client to cope with the data heterogeneity. Personalization can be done via multiple methods such as layer personalization [2, 13, 14] , regularization term [15, 16] , fine-tuning [17, 18] , and knowledge distillation [19, 20] . Although personalization can effectively tackle non-iid datasets, it requires extra computation [15, 17, 18, 19, 20] for resource-constrained clients, and its communication overhead can also be high due to large model sizes. As such, computation and communication costs are still one of the major challenges for personalized FL. This has motivated a line of research in training sparse personalized models [7, 9, 21, 22, 23, 24] by allowing clients to locally prune a global model during training.To reduce the computation and communication overhead of complex ML models during training, the idea of embedding FL algorithms with pruning has recently emerged. In [5, 6, 7, 8, 9, 21, 22, 23, 24, 25, 26, 27] , the clients train sparse models and only communicate sparse model parameters to reduce computation and communication overhead. To improve the aggregation phase with sparse models, the works in [5, 8, 21] perform averaging only between overlapping parameters to avoid information dilution by excluding zero value parameters. The work in [6] obtained a sparse model by selecting a particular client to prune an initial dense model and then performed training in a similar way to FedAvg. In [25] , the authors presented binary masks adjustment strategy to improve the performance of sparse models and communication efficiency. The work in [26] progressively pruned a dense model for sparsification and analyzed its convergence. Similarly, in the FL solutions of [9, 21] , the clients train and communicate personalized sparse local models while iteratively pruning a dense model. In [7, 22] , the clients optimized personalized sparse models by exchanging lottery tickets [28] at every communication round. The work in [24] obtained personalized sparse models by l 1 norms constraints and the correlation between local and global models. The FL framework of [23] allows clients to train personalized sparse models in a decentralized setting without a central server. Although these works adopted sparse models during training, if models are not sparse enough, the computation and communication costs can remain high due to the large model sizes. Moreover, trained models often have a poor performance at high sparsity. In [29, 30] , binary masks are communicated and optimized by training auxiliary variables while freezing model parameters. However, the work in [29] approximated binarization step using a sigmoid function during forward propagation. In [30] , the downlink communication costs remained the same as that of FedAvg. In [10, 11, 31] , clients perform neural-architecture-search by training multiple models to find optimized and sparse models to improve computational and memory efficiency at inference phase. However, in practice, clients often have limited resources to support the computationally intensive architecture search process. Therefore, most prior works either required extensive computational resources to find sparse models, or achieved a poor performance for high model sparsity. Different from prior works, in the proposed SpaFL framework, we do not incur more computational resources to find sparse models nor do we jeopardize the performance for sparsification. 
spafl algorithm: 
 In this section, we first formulate our FL problem and present SpaFL to solve the proposed problem with low computation and communication costs.We now formulate our main problem. We aim to optimize each client's model parameters and sparse model structures jointly in a personalized FL setting. This can be formulated as the following bi-level optimization problem:EQUATIONwhere p k ∈ R d is a binary mask corresponding to a personalized sparse model structure of client k, F k (•) is a empirical risk associated with local data of client k, L is a loss function, D k is the number of data samples, {x, y} is an input-label pair, w k captures the model parameters, and ⊙ is the Hadamard product. If the element of p k is zero, then the corresponding parameter of w k will be pruned. Our goal is to obtain the optimal w k and p k for each client in order to reduce the computation and communication overhead during training while being robust to the data heterogeneity among clients. However, solving (1) is not trivial because w k and p k are highly correlated [11] .Moreover, personalization and model sparsification should be achieved in a resource-efficient manner because most clients (e.g., mobile or IoT devices) and communication systems do not have enough computation and communication resources.We now describe the proposed algorithm, SpaFL, that can solve (1) while maintaining communicationefficiency with low computational cost. In SpaFL, every client jointly optimizes its personalized sparse model structure and model parameters with trainable thresholds, which can be used to prune the model parameters. To save communication resources, only thresholds will be aggregated at a server to generate global thresholds for the next round. Here, global thresholds can represent the aggregated parameter importance of clients. Hence, at the beginning of each round, every client extracts the aggregated parameter importance from the global thresholds so as to update its model parameters. The overall algorithm is illustrated in Fig 1 . and summarized in Algorithm 1.To optimize the binary mask p k in (1), inspired from [32] , we define a trainable threshold for each neuron in fully-connected layers or for each filter in convolutional layers. The neural network of client k will consist of L layers as{W 1 k , . . . , W L k }. For parameters W l k ∈ R n l out ×n lin in a fully-connected layer l, we define trainable thresholdsτ l ∈ R n l out . If it is a convolutional layer W l k ∈ R n l out ×c l in ×k l ×h l, where c l in is the number of input channels and k l × h l are the kernel sizes, we can changeW l k as W l k ∈ R n l out ×n l in with n l in = c l in × k l × h l .Similarly, we can define the corresponding thresholds τ l ∈ R n l out . Then, for each client k, we define a set of total thresholds τ = {τ 1 , . . . , τ L }. Note that the number of these additional thresholds will be at most 1% of the number of model parameters d. Moreover, to further improve the communication efficiency, only these thresholds will be communicated between clients and the server.We prune a parameter if its magnitude is smaller than its connected neuron/filter's threshold. Then, we can obtain a binary mask p l k for W l k , as followsEQUATIONwhere S(•) is a unit step function. Hence, we can obtain the binary masks {p 1 k , . . . , p L k } by performing (2) at each layer. To facilitate the pruning, we constrain the parameters and thresholds to be within [-1, 1] and [0, 1], respectively [32] . For simplicity, we unroll {W 1 k , . . . , W L k } and {p 1 k , . . . , p L k } to w k ∈ R d and p k ∈ R d , respectively as done in [33] .At each round, a server samples a set of clients S t such that |S t | = K for local training. For given global thresholds τ (t) at round t, client k ∈ S t generates a binary mask p k (t; τ (t)) using (2). Subsequently, it obtains the sparse model wk (t) = w k (t) ⊙ p k (t; τ (t)). To improve communication efficiency, each client performs E epochs with p k (t; τ (t)). In particular, for E > 1, client k ∈ S t performs E -1 epochs to its sparse model using mini-batch stochastic gradient as follows:EQUATIONwhere g k ( we k (t)) = ∇ we k F k ( we k (t), τ (t); ξ e k (t)) with a mini-batch ξ and η(t) is a learning rate for the model parameters. During E -1 epochs, the model parameters can adapt to the new binary mask p k (t; τ (t)), which is generated using the current global thresholds τ (t), by recovering from pruning-induced noise [34] . Then, in the last epoch e = E -1, client k updates the received global thresholds τ (t) via backpropagation. Client k first calculates the following sparsity regularization term R(t) = L l=1 n l out i=1 exp(-τ i ). Then, the loss function at the last epoch e = E -1 will be:EQUATIONwhere 0 ≤ α ≤ 1 is the coefficient that controls R(t). From (4), we can give thresholds τ (t) performance feedback on the current sparse model while also progressively increasing τ (t) through the sparsity regularization term R(t) [32] . Hence, prematurely pruned parameters can be recovered since thresholds and parameters are jointly updated while progressively enforcing sparsity. From (4), client k then updates the received global thresholds τ (t) via backpropagation as followsEQUATIONwhereEQUATIONFor simplicity, we can vectorize (5) as follows:EQUATIONwhereh k ( wE-1 k (t)) = ∇ τ F k ( wE-1 k (t), τ (t); ξ E-1 k (t)). When we calculate the gradients of thresholds τ (t), we use the identity straight-through estimator [35] to approximate the gradient of the step functions in (2) .From ( 6), we can see that threshold τ l k,i corresponds to the importance of its connected parametersw l k,ij , 1 ≤ j ≤ n l in , in the sparse model wk (t). This is because the importance of a parameter w l ij can be estimated by [36] EQUATIONwhere F (w, τ ; w l ij = 0) is the loss function when w l ij is masked and the approximation is obtained from the first Taylor expansion at w l ij = 0. Therefore, if connected parameters were important, the sign of (8) of those parameters will be negative, and the corresponding threshold will decrease as in (6) . Otherwise, the threshold will be increased to enforce sparsity. Hence, prematurely pruned parameters will be automatically recovered via a joint optimization of τ and w.After local training, each client k ∈ S t , transmits the updated thresholds τ k (t) to the server. Here, the communication overhead will be less than one percent of that of transmitting the entire parameters. Subsequently, the server aggregates the received thresholds and generates new global thresholds for the next round, i.e.,EQUATIONSince thresholds represent the importance of the connected parameters at each neuron/filter, clients can learn how to prune their parameters from the global thresholds. Moreover, the difference between two consecutive global thresholds ∆τ (t) = τ (t + 1) -τ (t) captures the history of aggregated parameter importance, which can be further used to improve model performance. For instance, from (8), if ∆τ l i (t) < 0, then the parameters connected to threshold i in layer l were globally important. If ∆τ l i (t) ≥ 0, then the connected parameters were globally less important. Hence, from ∆τ (t), clients can deduce which parameter is globally important or not and further update their model parameters. After generating new global thresholds τ (t + 1), the server broadcasts τ (t + 1) to client k ∈ N , and then clients calculate ∆τ (t) = τ (t + 1) -τ (t).We now present how clients can update their model parameters from ∆τ (t). For given ∆τ (t), we need to decide on the: 1) update direction and 2) update amount. Clients can know the update direction of parameters by considering ∆τ (t) and the dominant sign of parameters connected to each threshold. For simplicity, assume that each parameter has a threshold. Then, the gradient of the thresholds in ( 6) can be rewritten as follows:EQUATIONThe gradient of the loss F k ( wk (t), τ (t)) with respect to the whole parameters w k (t) is given byEQUATIONFrom ( 10) and (11) , the gradient direction of a parameter w is opposite of that of its connected threshold if w > 0. Otherwise, both the threshold and the parameter have the same gradient direction. Hence, we can deduce the following: If w > 0, the gradient direction of w and the sign of ∆τ will have the same sign; otherwise, the gradient direction of w and the sign of ∆τ are opposite. Since in SpaFL each threshold has multiple connected parameters, we decide the update direction of connected parameters by finding the dominant sign among them. To this end, we simply add the connected parameters of each threshold. For instance, consider threshold i in layer l of client k, ifEQUATIONt) > 0, then the gradient direction of the connected parameters will be the same as the sign of ∆τ l i (t). Otherwise, it is the opposite of the sign of ∆τ l i (t). Thus, the update direction can be simply expressed with a XOR operation between the sign of ∆τ l i (t) and the sign of connected parameters sum. Next, we decide how much a parameter should be updated. From (10) and (11) , we can see that a threshold and a parameter have the same magnitude for their gradients. Hence, we simply divide ∆τ l i (t) by the number of connected parameters n l in . We finally provide the update equation using ∆τ (t) as followsEQUATIONwhere sign(•) is a sign function. This parameter update corresponds to line 7 in Algorithm 1. Note that this additional parameter update is not computationally intensive because it happens only once before local training. We also provide the number of used FLOPs during training with inclusion of this operation in Section 5. 
theoretical analysis of spafl: 
 We now present our analysis on the convergence of thresholds under the following commonly adopted assumptions [37] .EQUATIONAssumption 2. (Unbiased stochastic gradient) The stochastic gradient h k is an unbiased estimator of the gradient ∇ τ F k , respectively, for client k, ∀k, such thatEQUATIONThen, we can derive the convergence of thresholds in the following theorem.Theorem 1. For γ(t) = η(t)(1 -α(1-M η(t))2) and the largest number of parameters connected to a neuron or filter n max in > 0 in a given model, we haveEQUATIONProof is provided in the Supplementary document.From (6), thresholds τ (t) are updated using parameter gradients g k (t), k ∈ S t . We can expect that the thresholds will converge when parameters w k , ∀k, converge. We can see that the sparsity regularizer coefficient α impacts convergence. As α increases, we can quickly enforce more sparsity to the model. However, a very large α will damage the performance as γ(t) decreases in (??). We can also observed that the convergence depends on the difference between the received global thresholds τ (t) and the updated thresholds τ k (t). Hence, a very large change to the global thresholds will lead to a significantly different binary mask in the next round. Then, local training can be unstable as parameters have to adapt to the new mask. Therefore, from Theorem 1, we can capture the tradeoff between the computing cost and the learning performance in terms of α. 
experiments: 
 We now present experimental results to demonstrate the performance, computation costs and communication efficiency of SpaFL. Implementation details are provided in the Supplementary document.We conduct experiments on three image classification datasets: FMNIST [38] , CIFAR-10, and CIFAR-100 [39] datasets. To distribute datasets in a non-iid fashion, we use Dirichlet (0.2) for FMNIST and Dirichlet (0.1) for CIFAR-10 and CIFAR-100 datasets as done in [40] with N = 100 clients. We set the total communication round T = 500. At each round, we randomly sample K = 10 clients. Unless stated otherwise, we average all the results over at least 10 different random seeds. We also calculate the best accuracy by averaging each client's performance on its test dataset.For FMNIST dataset, we use the Lenet-5-Caffe and LSTM models. For the Lenet model, we set η(t) = 0.001, E = 3, α = 0.002, and a batch size to be 64. For the LSTM model, we use two LSTM layers with hidden size of 128, η(t) = 0.01, E = 3, α = 0.0003 and a batch size of 16. For CIFAR-10 dataset, we use a convolutional neural network (CNN) model with seven layers used in [41] with η(t) = 0.01, E = 3, α = 7.5 × 10 -5 , and a batch size of 16. We adopt the ResNet-8 model [42] for CIFAR-100 dataset with η(t) = 0. Table 2 : Performance of SpaFL and other baselines along with their used communication costs (Comm) and computation (FLOPs) resources during whole training with the LSTM model.We compare SpaFL with seven state of the art baselines that include both dense and sparse FL algorithms. In FedAVG [1] , every client trains a global dense model and communicates whole model parameters.LG-FedAvg [2] is a dense personalized FL algorithm, where clients learn local representations and share only small global layers for communication efficiency. PruneFL [6] learns a global sparse model after pruning a dense model at a particular client for initialization. Sub-FedAvg [9] is a scheme that learns a personalized sparse model for each client and iteratively performs pruning using validation dataset during training. LotteryFL [7] optimizes personalized sparse models by communicating lottery tickets of clients at every round. In FedDST [5] , clients learn a global sparse model and a global binary mask while communicating only their sparse model parameters and binary masks. FedSpa [24] trains personalized sparse models for clients while maintaining fixed model density during training. FedPM [30] trains and communicates a binary mask while freezing model parameters. For the sparse FL baselines, the target sparsity is set to 0.5 following the configurations in [5, 6, 7, 9, 24] .In Table 1 and Fig. 2 , we present the averaged accuracy, communication costs, number of FLOPs during training, and convergence rate for each algorithm. We consider all uplink and downlink communications to calculate the communication cost of each algorithm. We also provide the details of the FLOPs measure in the Supplementary document. We average the model densities of SpaFL when a model achieved the best accuracy during training. From these results, we observe that SpaFL outperforms all baselines while using the least amount of communication costs and number of FLOPs except FMNIST dataset. The achieved model densities are 5.36%, 7.57%, and 7.38%, for FMNIST, CIFAR-10, and CIFAR-100, respectively. We also observe that SpaFL uses less resources and performs better than Sub-FedAvg and LotteryFL, which also personalize both model parameters and sparse model structures. Although FedPM reduced uplink communication costs by communicating only binary masks, its downlink cost is the same as FedAvg. In SpaFL, since the clients and the server only exchange thresholds, we can significantly reduce the communication costs compared to baselines that exchange the subset of model parameters such as LG-FedAvg, Sub-FedAvg, LotteryFL, PruneFL, and FedDST. Hence, SpaFL can efficiently improve model performance with small computation and communication costs. In Table 2 , we also compare the performance and resource usage of SpaFL with baselines that considered LSTM models in their work. In Fig. 2 , we show the convergence rate of each algorithm. We can see that the accuracy of SpaFL decreases and then keeps increasing. The initial accuracy drop is from pruning while global thresholds are not trained enough. As thresholds keep being trained and communicated, clients learn how to prune their model, thereby gradually improving the performance even at high model sparsity.In Fig. 3 , we show the change of model density of SpaFL during training with a different sparsity coefficient α. From Fig. 3 , we can observe that the model density fluctuates at its low value. This is because we are jointly optimizing both model parameters and thresholds. Hence, prematurely pruned parameters will be recovered during training. In Fig. 3 and Table 3 , we observe the tradeoff between model performance and computational costs. As α increases, we can quickly enforce sparsity to the model. However, the performance can be low if we set α to be very large. This is because binary masks can change too quickly, thereby making training unstable. Therefore, Fig. 3 also corroborates our theoretical analysis in Theorem 1. 
conclusion: 
 In this paper, we have developed a communication-efficient FL framework SpaFL that allows clients to optimize both personalized model parameters and sparse structures with low computational costs.We have reduced computational overhead by performing pruning through trainable thresholds. To further reduce communication costs, we have communicated only thresholds between clients and a server. We have also presented the parameter update method that can extract parameter importance from global thresholds. Furthermore, we have provided theoretical insights on the convergence of thresholds and experimental results to demonstrate the resource-efficiency of SpaFL.Under review as a conference paper at ICLR 2024 
: 
 trainable thresholds. SpaFL communicates only the thresholds so as to learn how to prune and significantly save communication costs. Here, a trainable threshold is defined for each neuron/filter and is used to prune its connected parameters based on magnitude. Both parameters and thresholds are jointly optimized to enable the automatic sparsification of models while recovering prematurely pruned parameters during training. Therefore, SpaFL reduces computational overhead on resource constrained devices with minimal performance loss. To further save communication costs, only thresholds are communicated between clients and a server. Hence, clients can learn how to prune their model from global thresholds. Since parameters are not communicated, the clients' parameters and sparse model structures will remain personalized while only global thresholds are shared. We show that global thresholds can capture the aggregated parameter importance of clients. We further update the clients' model parameters by extracting aggregated parameter importance from global thresholds to improve performance. We analyze the convergence of thresholds while shedding light on the tradeoff between computational overhead and performance. We summarize our contributions as follows:• We propose a new communication-efficient FL framework called SpaFL, in which clients optimize their personalized model parameters and sparse model structures with low computing costs through trainable thresholds.• We show how SpaFL can significantly reduce communication overhead for both clients and the server by only exchanging thresholds, the number of which is less than two orders of magnitude smaller than the number of model parameters.• We prove the convergence of thresholds. Moreover, the impact of thresholds on the model performance is theoretically and experimentally analyzed.• Experimental results demonstrate the performance, computation costs, and communication efficiency of SpaFL compared with both dense and sparse baselines. For instance, the results show that SpaFL uses only 2.87% of communication and 14.7% of computation resources compared to a dense baseline FedAvg while improving accuracy. Additionally, SpaFL improves accuracy by 2.13% compared to a sparse personalized baseline while consuming only 25.55% of this baseline's communication resources, and only 14.67% of its computing resources. 
 
 Appendix: We next present the detailed proof of Theorem 1. To facilitate the proof, we first restate the update rule of thresholds and used assumptions in the main paper. In SpaFL, thresholds τ (t) are updated as followsEQUATION)We then prove Theorem 1 under the following commonly adopted assumption [37] .EQUATION)Assumption 4. (Unbiased stochastic gradient) The stochastic gradient h k is an unbiased estimator of the gradient ∇ τ F k , respectively, for client k, ∀k, such thatEQUATIONWe first consider the case in which global thresholds converge. We have the following update rule for global thresholds asEQUATIONWe take the expectation over the randomness in client scheduling and stochastic gradients as followsEQUATIONHence, clearly τ will eventually converge ifEQUATION)|| 2 converges. We next show that this conditional statement holds in our SpaFL framework.From the M -smoothness of the loss function in Assumption 3, we haveEQUATIONTo facilitate the analysis, we first derive τ k (t) -τ (t) as belowEQUATIONThen, we can change (21) as followsEQUATION)We next take the expectation to the above inequality and use Assumption 4 as belowEQUATION)We first bound A using ⟨a, b⟩ ≤ ||a|| 2 +||b|| 2 2 as belowEQUATION)We now further bound B using ( 16) asEQUATIONwhere n max in is the largest number of parameters connected to a neuron or filter in a given model, (a) is from |w| ≤ 1 in Section 3.2.1, and the last inequality is from the M -smoothness of F k . By combining A and B with taking expectation, we haveEQUATIONBy arranging the above inequality, we haveEQUATIONwhere γ(t) = η(t)(1 -α(1-M η(t))2). We now further bound (A) in (28) . From Assumption 3, we have the followingEQUATION)Based on (29) , we can bound (28) as belowEQUATION)From (30) , we can bound the averaged aggregated gradients with respect to thresholds as belowEQUATION)By summing the above inequality from t = 0 to t = T -1, we can obtain the result of Theorem 1.In Theorem 1, we have the loss function F k , which is not assumed to be bounded in our analysis.Although loss functions can generally be a large value, we constrained model parameters w and thresholds τ to be within [-1, 1] and [0, 1], respectively. Therefore, if input-output pairs x, y are drawn from some bounded domain, the loss function F k can be obviously bounded by a finite real value. However, since inputs can be sampled from unbounded domain, we provide Corollary below to mitigate this issue. Corollary 1. Assume that there existsG ≥ 0 such that E||g k (w k )|| 2 ≤ G 2 , ∀k. For γ(t) = η(t)(1 -α(1-M η(t))2) and the largest number of parameters connected to a neuron or filter n max in > 0 in a given model, we haveEQUATIONProof. We start from bounding the term B in eq. ( 24) as followsEQUATION)By following eq. ( 30) and ( 31), we can derive the convergence bound.We calculate the communication cost of SpaFL considering both uplink and downlink communications. At each round t, sampled clients transmit their updated thresholds to the server. Hence, the uplink communication costs can be given byEQUATIONwhere τ num is the number of thresholds of a given model. In downlink, the server broadcasts the updated global threshold to every client. Hence, the downlink communication costs can be given as belowEQUATIONTherefore, total communication costs can be given by T × (Comm Up + Comm down ).We calculate the number of FLOPs during training using the framework introduced in [43] . We consider a convolutional layer with an input tensor X ∈ R × N × (C × R × S) × F × H × W , where ρ = ||p||0 ||W ||0is the density of the parameter matrix W . For the backpropagation, we calculate input gradient and parameter gradient. We can calculate the number of FLOPs for the input gradient by convolving the output gradient with the parameter matrix W . Hence, it can approximated to N ×(F ×R×S)×C ×X ×Y . Since we used a sparse model during the forward propagation, this can be reduced to ρ × N × (F × R × S) × C × X × Y . For the parameter gradient, we can calculate it by convolving the input activation with the output gradient. Therefore, this can be approximately given by (N × X × Y ) × F × C × R × S. Here, we did not block the gradients of input activations. Thus, the FLOPs during the backward propagation can be given by(1 + ρ)(N × F × R × S × C × X × Y ).Hence, the computational overhead of backpropagation is approximately (1 + ρ) times that of the forward propagation.For a fully connected layer with input tensor X ∈ R N ×X and parameter tensor W ∈ R X×Y , the input tensor X is multiplied with W during the forward propagation. Hence, with the density of W , we can calculate the number of FLOPs for the forward propagation as ρ × N × X × Y . In backpropagation, we can obtain the gradient of loss by calculating the input gradient and parameter gradient in a similar manner with convolutional layers. Hence, the number of FLOPs during backpropagation can by given by (1 + ρ) × N × X × Y .For an LSTM layer with input size I and hidden size D, we can derive the number of FLOPs of the forward propagation as 4(I × D × ρ 1 + D × D × ρ 2 ), where ρ 1 and ρ 2 corresponds to the density of the corresponding weight matrices [44] . Note that an LSTM layer usually has three gates and one memory cell. We approximated the number of FLOPs of backpropagation as(1 + ρ1+ρ2 2 ) × (I × D × ρ 1 + D × D × ρ 2 )as done in linear and convolutional layers previously. The number of FLOPs for model parameter update will be the same as the number of model parameters d. For, thresholds update, it can be given by ||τ || 0 + d because we perform multiplication and addition using the connected parameters of each threshold. We also consider the number of FLOPs to perform line 7 in Algorithm 1 for updating the local models from global thresholds. Every client first has to decide update directions by doing summation of connected parameters at each neuron/filter (sum operation). Then, they update their local models using the received global thresholds (sum and multiply operations). This corresponds to 1.5 × d FLOPs, where d is the number of model parameters. Then, the total number of FLOPs during one local epoch at round t can be approximately given by SpaFL can be readily extended to perform structured pruning. Specifically, we can prune entire parameters connected to a filter/neuron. In SpaFL, we defined a threshold for each filter/neuron and pruned each connected parameter if its magnitude is smaller than that of the threshold. This approach leads to unstructured sparsity. We can endow SpaFL with structured sparsity by calculating the average magnitude of all parameters connected to their filter/neuron and pruning all of them if their average magnitude is smaller than their threshold. Hence, SpaFL can be extended to do filter/neuron-wise pruning. We provide the performance of SpaFL with structured sparsity below. The achieved model densities are 50.6%, 50%, and 49.3% for FMNIST, CIFAR-10, and CIFAR-100, respectively. The target density of the baselines is set to be 50%. Since it is well known that filter/neuron-wise pruning is less flexible than unstructured pruning, the results show accuracy loss compared to the original SpaFL. However, we can see that the structured SpaFL still outperforms unstructured baselines at the same model density 50%.FLOP(t) = L l=1 (2 + ρ l (t))N × (C l × R l × S l ) × F l × H l × W l × 1{layer l == conv} + (2 + ρ l (t)) × N × X l × Y l ×In this subsection, we investigate the impact of different sparsity regulaizer R(t) for updating thresholds used in (4). We chose R(t) = L l=1 n l out i=1 exp(-τ i ) because an exponential function goes to zero asymptotically as thresholds increase. This is reasonable because it can penalize low thresholds without making them become extremely large. We can also use a similar regularizer that can give penalty for low thresholds without encouraging them to become a too large value. For example, a regularizer R ′ (t) = 1 ||τ (t)+0.1 satisfies such property. Below, we provide an empirical comparison between SpaFL and the baseline that uses R ′ (t) in Table 7 .For the regularizer R(t), the achieved model densities are 5.36%,7.57% and 7.38% for FMNIST, CIFAR-10, and CIFAR-100, respectively. Meanwhile, with R ′ (t), the model densities are 4.85%, 5.60%, and 8.01% for FMNIST, CIFAR-10, and CIFAR-100, respectively. We can see that both regularizers work well with SpaFL. Hence, SpaFL is compatible with any regularizers that can penalize the magnitude of thresholds asymptotically. 


Decision: This paper introduces SpaFL: a federated learning framework for enhancing communication efficiency and minimizing computational overhead by optimizing sparse model structures. This goal is achieved by defining a trainable threshold which leads to structured sparsity. Since the server and clients exchange thresholds only, SpaFL is able to reduce communication overhead. Furthermore, these trainable thresholds lead to lower computation overhead, too.

All reviewers see the paper favorably. In particular, all give acceptance-level scores, and list several compelling strengths, including:

- Communicating the threshold instead of the model weights is an interesting idea
- Preliminary empirical results on popular benchmarks seem to have shown that this idea indeed lowers communication overhead
- The paper is well-written and easy to follow

The several weaknesses that were pointed out by the reviewers were mostly addressed in the rebuttal and ensuing discussion. I fully expect these developments, clarifications and explanations to be included in the camera-ready version of the paper.

Reviews: review: paper_summary: paper_summary: This paper suggests communicating the threshold instead of the model parameters in federated learning. Through empirical validations on popular benchmarks, the proposed method, SpaFL, is shown to have lower computational overhead and achieve relatively good results.
strength_weakness: strength_weakness: 1. The main figure, Figure 1, is not informative and very confusing. I would expect to see a high-level explanation of how the threshold for each client is selected, whether they are the same for each layer, and how the selected thresholds are combined. At this moment, I can only see that all local models have the same threshold and they are directly combined at the server. Due to the misalignment between the main figure and the illustration in section 3.1, the necessary details to fully understand the main idea of this paper are not easy to follow. The writing in this section could also be improved.

2. In Theorem 1, the authors provided a generalization bound. However, the intuition and analysis for this bound are missing. At this moment, I do not have a good sense of, in a practical or special case, how large the right-hand side probability of Equation 14 is? How large must the minimal training data size $D$ be? Is it practical? Also, this bound only measures the distance between the empirical risk and the expected risk; how can this be regarded as a generalization bound? To my understanding, a generalization bound should be between the training set and the testing set. Furthermore, could you please provide the convergence rate and communication complexity of the proposed method? These are very important for the theoretical analysis of FL.

3. The baselines are outdated, such as Fjord (NeurIPS21) and HeteroFL (ICLR21). Can you please compare more recent methods, which can be easily found, such as [FedPAC](https://arxiv.org/abs/2306.11867) and [FedCR](https://proceedings.mlr.press/v202/zhang23w.html)? I would expect to see more comparisons with newer methods. They do not need to be the ones mentioned above, but I believe it is critical to determine whether the proposed method is still interesting to explore at this moment.
questions: questions: I would expect a README file to better understand your code efficiently.
limitations: limitations: The major limitations include the unclear presentation of the proposed method, lack of analysis of the generalization bound, absence of convergence and communication efficiency analysis, and the use of outdated empirical baselines.
paper_summary: paper_summary: This paper introduces SpaFL, a federated learning framework that enhances communication efficiency and minimizes computational overhead by optimizing sparse model structures. They achieve this goal by defining a trainable threshold which leads to structured sparsity. Since the server and clients only exchange thresholds, SpaFL is able to reduce communication overhead. Furthermore, these trainable thresholds also lead to low computing costs.
strength_weakness: strength_weakness: 1. I recommend the authors to do additional experiments with NLP tasks.

2. In Figure 8 of the FedPM paper, the accuracy achieved by FedPM is higher than what is reported in this paper. Could you clarify the reason for this discrepancy? Is it due to differences in experimental setup and hyperparameters? If so, could you ensure that all conditions are consistent for the CIFAR-100 experiments and compare your method to the original accuracy reported in the FedPM paper?
questions: questions: 1. If you apply this method to the transformer based models, how can you define trainable thresholds? To be specific, you gave an explanation for a convolutional layer in section 3.1, and can we just similarly define for self attention modules?
limitations: limitations: N/A
paper_summary: paper_summary: The authors consider distributed training of sparse models, by optimizing over the thresholds used to prune the models.
strength_weakness: strength_weakness: I would be interested in a discussion and comparison with other approaches that train sparse models, such as
* Meinhardt et al. "Prune at the Clients, Not the Server: Accelerated Sparse Training in Federated Learning," preprint arXiv:2405.20623, 2024.  
* Yi et al. "FedComLoc: Communication-Efficient Distributed Training of Sparse and Quantized Models," preprint arXiv:2403.09904, 2024. 

and references therein.

The experiments should investigate more the influence of the level of heterogeneity, because this is what makes it difficult to identify the sparsity pattern.
questions: questions: No question
limitations: limitations: The technical limitations are discussed at the end of the paper.

Comments: title: Thanks!
raise to 6
comment: We appreciate the reviewer for the constructive comments. We will update the convergence analysis, communication costs, and the comparisons with more recent baselines in the revision
Thanks for your response. I will keep my current score.
Thanks to all reviewers for engaging with the authors!!

Reviewer 28Jj: I'd appreciate if you could read the latest response of the authors and reply to it.

Thanks!

AC
We appreciate the reviewer for the constructive comment. We will update the revised version with newly added baselines and sparsity patterns with decreasing data heterogeneity.
We appreciate the reviewer for the constructive comment. We will update the revised version with the proposed method for transformer models.
We appreciate the reviewer for the constructive comments. We provide our response to each comment below. 

>Q1:  Assumption 3 regarding constant bounded gradient is quite strong and may not be feasible in the context of FL

**A1:** In our Theorem 1, we derived the generalization guarantee through differential privacy. Specifically, we showed that the proposed algorithm satisfies differential privacy at each communication round by following [1]. Although our work does not focus on privacy, we will discuss about privacy guarantees in the revised version, 

>Q2: The suggested theoretical comparison with existing methods is still missing.

**A2**: We compare our convergence analysis with [2], which provided the convergence rate of FL with arbitrary time-varying pruned models. Specifically, clients train pruned models, whose binary masks can change over time. and communicated pruned models with the server to generate a global model. The authors in [2] also made the same Assumption 1 and 3. Then, they further assumed that $||\theta_q - m_q \odot \theta_{q, n}||^2 \leq \delta^2 ||\theta_q||^2$, where $\theta_q$ is a global model at round $q$, $m_q$ is a binary mask, $0 \leq \delta < 1$ is pruning induced noise, and $\theta_{q, n}$ is the local model of client $n$. According to the Theorem in [2], they provide

$\frac{1}{T} \sum_{t=1}^{T} \mathbb{E}  || \nabla F(\theta_q)||^2 \leq \frac{G_0}{\sqrt{T}} + \frac{V_0}{\sqrt{T}}  + \frac{H_0}{T} + \frac{I_0}{\Gamma\_{min}} \sum_{t=1}^{T} \mathbb{E} ||\theta_q||^2$, where $\Gamma\_{min}$  measures the minimum occurrence of the parameter in the local models in all rounds.

Now we compare our convergence rate with the one from [2]. Although both algorithm converges to a stationary point with the rate of $1/\sqrt{T}$, we note that [2] has a non-vanishing term $\frac{\delta^2 I_0}{\Gamma\_{min}} \sum_{t=1}^{T} \mathbb{E} ||\theta_q||^2$ due to the noise from pruning. We also have similar terms such as  $\frac{\sum\_{t=0}^{T-1}  4\alpha || \exp(-\tau(t) ) ||^2}{T(1-\alpha)}$ due to the regularization. The main difference is that the convergence analysis of [2] cannot effectively capture the impact of sparse models due to uncontrollable and theoretical value $\delta$. Meanwhile, in our analysis, we can effectively show the impact of the regularization coefficient $\alpha$ on the convergence rate. 

>Q3: The calculation of uplink and downlink communication costs remains unclear. 

**A3:** We calculated the uplink and downlink communication costs from physically communicated amount of data between clients and the server during the fixed number of rounds. Specifically, each parameter can be represented with 32bit. For the uplink communication cost, we measured how many parameters are transmitted from clients to the server and then multiplied the number of transmitted parameters by 32 [bit] to derive the cost. We do a similar calculation for the downlink by measuring the number of parameters sent to clients. As the reviewer mentioned, we note that calculating the communication cost from the convergence rate can provide better understanding of the communication complexity as done in [3]. We will clarify how we calculated the up/down link communication costs and also provide the result by following [3] in the revised version. 

>Q4: While Fjord remains a popular method for comparison, I look forward to seeing comparisons with more recent approaches. 

**A4**: In the revised version, we will include the comparison with [4] and [5], which studied structured sparsity in FL.

[1] Rong Dai, Li Shen, Fengxiang He, Xinmei Tian, and Dacheng Tao. Dispfl: Towards communication-efficient personalized federated learning via decentralized sparse training. In
International Conference on Machine Learning, pages 4587–4604. PMLR, 2022

[2] Zhou, Hanhan, et al. "Every parameter matters: Ensuring the convergence of federated learning with dynamic heterogeneous models reduction." Advances in Neural Information Processing Systems 36 (2024).

[3]  Malinovsky, Grigory, Kai Yi, and Peter Richtárik. "Variance reduced proxskip: Algorithm, theory and application to federated learning." Advances in Neural Information Processing Systems 35 (2022): 15176-15189.

[4] Dongping Liao, Xitong Gao, Yiren Zhao, and Cheng-Zhong Xu. Adaptive channel sparsity for federated learning under system heterogeneity. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20432–20441, 2023.

[5] Chen, Daoyuan, et al. "Efficient personalized federated learning via sparse model-adaptation." International Conference on Machine Learning. PMLR, 2023.
Thank you for your follow-up response. I appreciate the additional comparison with FedP3. I would suggest that the authors incorporate the points we've discussed into the paper, as this could significantly enhance both the clarity and contribution of the work. I also want to acknowledge the effort put into the rebuttal and the follow-up response.

I would like to highlight a few important points:

a) Assumption 3 regarding constant bounded gradient is quite strong and may not be feasible in the context of FL. It only seems reasonable when considering differential privacy to achieve privacy guarantees (perhaps future work on differential privacy could also help alleviate this.) This could be a potential limitation of your work, and it would be valuable to see a proper discussion of this in your future revisions. Additionally, directly assigning a complex term, such as A, B, etc., to the final convergence rate—especially when it depends on many variable factors, including the function itself—is not a reasonable way unless you have thoroughly examined its characteristics.

b) The suggested theoretical comparison with existing methods is still missing. This will likely require substantial effort to carefully examine each variable in your results to ensure a fair comparison. Including this will greatly improve the soundness of the work.

c) The calculation of uplink and downlink communication costs remains unclear. For example, you might follow Eqn (15) in [1] to provide an explicit computational analysis of the cost, which would be more convincing in demonstrating the potential benefits.

d) While Fjord remains a popular method for comparison, I look forward to seeing comparisons with more recent approaches. I’m glad you included a comparison with FedP3 in the rebuttal, even though your focuses differ. If possible, I suggest incorporating additional comparisons with recent baselines, which would more convincingly demonstrate the merits of this work.

Overall, while I believe there is substantial room for improvement in terms of clarity and completeness, I see this work as interesting and promising. I have decided to raise my score.

[1] Malinovsky, Grigory, Kai Yi, and Peter Richtárik. "Variance reduced proxskip: Algorithm, theory and application to federated learning." Advances in Neural Information Processing Systems 35 (2022): 15176-15189.
Thank you for the constructive suggestions. We provide our response to each comment. 

>Q1: Could you please provide the explicit convergence rate and communication cost of your method, denoted in big-O notation? When compared with existing work, what insights can be drawn regarding the theory?

To derive the explicit convergence rate in big-O notation, we follow the steps in [1] and [2]. We first assume $\eta = \frac{1}{\sqrt{T}}$ and $0 \leq \alpha < 1$. Then, we can bound $\frac{1}{\gamma(t)} \leq \frac{2}{\eta(t) (1 - \alpha)}$. By replacing $\eta(t)$ and $\gamma(t)$ with their assumed value and bound into the above convergence rate, we have the following bound.

$\frac{1}{NT} \sum_{t=0}^{T-1} \mathbb{E} || \sum_{k=1}^N \nabla_{\tau} F_k (\tilde{w}_k(t), \tau(t)) ||^2 $ 
$\leq \frac{A}{ \sqrt{T} (1-\alpha)} + \frac{ B }{ T (1-\alpha)} + \frac{C}{\sqrt{T}} + \frac{D}{\sqrt{T}}$, where


$A = \sum_{t=0}^{T-1}  \sum_{k=1}^N \frac{\mathbb{E} ||  \nabla_{\tau} F_k (\tilde{w}\_k(t), \tau(t) ) - \nabla_{\tau_k} F_k ( \tilde{w}_k(t), \tau(t) ) ||^2}{MN}$,   $B =  \sum\_{t=0}^{T-1}  4\alpha || \exp(-\tau(t) ) ||^2$ 

$ C =  M^2 n\_{in}^{max} G^2/2 $,  $ D = \sum_{t=0}^{T-1}  \sum_{k=1}^N \frac{\mathbb ||\tau(t) - \tau\_k(t)||^2 }{N}$.

Therefore, the Big-O notation of the convergence rate becomes $\mathcal{O}(\frac{A}{ \sqrt{T} (1-\alpha)}) + \mathcal{O} ( \frac{ B }{ T (1-\alpha)}) + \mathcal{O}(\frac{C}{\sqrt{T}} ) + \mathcal{O} (\frac{D}{\sqrt{T}})$.

To derive the communication cost/complexity [1], we need to bound the left-hand side in the form of $\frac{1}{NT} \sum_{t=0}^{T-1} \mathbb{E} || \sum_{k=1}^N \nabla_{\tau} F_k (\tilde{w}_k(t), \tau(t)) ||^2 \leq \epsilon$. To this end, we further bound the convergence rate using the dominant term $1/\sqrt{T}(1-\alpha)$ on the right-hand side as follows

$\frac{1}{NT} \sum_{t=0}^{T-1} \mathbb{E} || \sum_{k=1}^N \nabla_{\tau} F_k (\tilde{w}_k(t), \tau(t)) ||^2 \leq \frac{D}{\sqrt{T}(1-\alpha)}$ , where $D = A +B +C +D$.  Then, to achieve $\epsilon$ stationary point, the communication complexity will be $\mathcal{O}(\frac{1}{\epsilon^2 ( 1-\alpha)^2})$. Since we only communicate thresholds at each round, the communication cost becomes $\mathcal{O} (\frac{ d\_{\tau} }{\epsilon^2 ( 1-\alpha)^2} )$, where $d\_{\tau}$ is the number of thresholds in a given model. We note that $d\_{\tau}$ is usually less than 1% of the number of parameters. 

From the above analysis, we can see that our convergence rate and communication complexity depends on the regularizer coefficient $\alpha$, which we used to enforce sparsity to models when updating thresholds. As $\alpha$ increases, it damages both the convergence rate and communication complexity. In the experiments, we used small values of $\alpha$ such as 0.001 or 0.00005 to balance the model sparsity and convergence rate (performance). Hence, our analysis shows how model sparsity can have impact on the convergence rate. 

>Q2: Can you elaborate on how the uplink, downlink, and total costs are computed? 

For the uplink, we measured the amount of parameters that clients sent to the server during the training. For the downlink, we measured the amount of parameters that the server sent to the clients during the training. The total cost is the sum of 'uplink cost' and 'downlink cost'. Then, we converted the number of communicated parameters to Gbit metric. 

>Q3: Why is the uplink always the same for FedP3 across different sparsity ratios?

When calculating the uplink costs, we followed the implementation of FedP3. In their implementation, clients receive two sets of layers. The first set is a set of dense layers and another set is a set of pruned layers. According to their implementation, clients send the first set to the server, and this is why the uplink costs remain the same across different pruning ratios. In case, if we apply  pruning to the first set, then we have the following table, 

| Algorithm     | Uplink costs (Gbit) | Downlink costs (Gbit)| Total costs (Gbit) |
| -------- | ------- |------- |------- |
| SpaFL   | 2.3040 | 2.3040| 4.6080| 
| FedP3 (0.5)    | 2.1923 | 2678.22| 2680.41 |
| FedP3 (0.7)    |3.0697 | 3749.50| 3752.57 | 
| FedP3 (0.9)   | 3.9467 | 4824.73| 4828.67 |
|FedAvg|5361| 5361 |10712|

We can still see that the our total costs are still smaller than that of FedP3.

[1] Yi, Kai, et al. "FedP3: Federated Personalized and Privacy-friendly Network Pruning under Model Heterogeneity." arXiv preprint arXiv:2404.09816 (2024).

[2] Yang, Haibo, et al. "Anarchic federated learning." International Conference on Machine Learning. PMLR, 2022.
Thank you for your response and patience with my questions. I have two requests:

1. Could you please provide the explicit convergence rate and communication cost of your method, denoted in big-O notation? 
When compared with existing work, what insights can be drawn regarding the theory?

2. Can you elaborate on how the uplink, downlink, and total costs are computed? Why is the uplink always the same for FedP3 across different sparsity ratios?
We next discuss the baseline comparison. We first provide the comparison between FedP3 and our method by changing the global pruning rate from 0.5 and 0.9. We measure the uplink and downlink communication costs of each method. For FedP3, we used the OPU2 method with the ordered dropout pruning as our paper studies structured sparsity. 

| Algorithm    | Accuracy (CIFAR-100) | Uplink costs (Gbit) | Downlink costs (Gbit)| Total costs (Gbit) |
| -------- | ------- |------- |------- |------- |
| SpaFL  | $40.80\pm0.54$    | 2.3040 | 2.3040| 4.6080| 
| FedP3 (0.5) | $37.73 \pm 0.42$     | 4.3853 | 2678.22| 2682.60 |
| FedP3 (0.7)    | $41.12 \pm 0.35$    |4.3853 | 3749.50| 3753.88 | 
| FedP3 (0.9)    | $44.42 \pm  0.37$  | 4.3853 | 4824.73| 4825.17 |
|FedAvg|  $35.51\pm0.10$ |5361| 5361 |10712|

Our method can achieve 35.38% of model density. We can see that our method outperforms FedP3 in terms of performance and communication costs at low model density (0.5). Although FedP3 performs better at high model density (0.7, 0.9), we note that our work focuses on optimizing sparse models with very low communication costs. We can see that the total communication costs of SpaFL is significantly less than FepdP3 (0.5). This means that SpaFL is more amenable to deployment on systems with limited resources. We also note that we only communicate trainable thresholds instead of doing full-layer aggregation of parameters.

For the baselines, please note that Fjord (2021) and HeteroFL (2021) are still widely used for comparing the performance structured sparsity in FL (e.g., see [1] and [2]). While [2] studied personalized structured sparse models, [1] focused on global structured sparsity models. We will add [1] as a baseline in the revised version. 

[1] Dongping Liao, Xitong Gao, Yiren Zhao, and Cheng-Zhong Xu. Adaptive channel sparsity for
federated learning under system heterogeneity. In Proceedings of the IEEE/CVF Conference on
Computer Vision and Pattern Recognition, pages 20432–20441, 2023.

[2] Chen, Daoyuan, et al. "Efficient personalized federated learning via sparse model-adaptation." International Conference on Machine Learning. PMLR, 2023.
Thank you for the constructive feedbacks. We first provide the convergence analysis. For the derivation, we need few conventional theoretical assumptions following [1] [2]. 

**A1**: Loss function $F_k(\cdot)$ is $M$ smooth for $\tau$ and client $k, \forall k$.

**A2**: The stochastic gradient $h_k$ is an unbiased estimator of the gradient $\nabla_{\tau} F_k$ for client $k, \forall k$ such that $\mathbb{E} h_k(w_k) = \nabla_{\tau} F_k(w_k, \tau)$

**A3**: There exists $G \geq 0$ such that $\mathbb{E}||g_k(w_k)||^2 \leq G^2, \forall k$.

In SpaFL, both parameters and thresholds are simultaneously updated, but only thresholds are communicated between clients and the server. Following the same notations in the paper, we assume $N$ number of clients , learning rate $\eta$, $T$ number of communication rounds, sparsity regularizer coefficient $\alpha$. and the thresholds update rule $\tau\_k(t) \leftarrow \tau(t) - \eta(t) h\_k(\tilde{w}\_k(t)) + \alpha \exp(-\tau(t))$. Then, we have the following convergence rate.

For $\gamma(t) = \eta(t) (1 - \frac{\alpha (1 - M \eta(t) }{2} )$ and the largest number of parameters connected to a filter/neuron $n_{in}^{max} >0$ in a given model, we have

$\frac{1}{NT} \sum_{t=0}^{T-1} \mathbb{E} || \sum_{k=1}^N \nabla_{\tau} F_k (\tilde{w}_k(t), \tau(t)) ||^2 $ 

$\leq \sum_{t=0}^{T-1}  \sum_{k=1}^N \frac{   \mathbb{E} ||  \nabla_{\tau} F_k (\tilde{w}\_k(t), \tau(t) ) - \nabla_{\tau_k} F_k ( \tilde{w}_k(t), \tau(t) ) ||^2} {  {MNT\gamma(t)}}$ 
$+  \sum\_{t=0}^{T-1}  \frac{ 2\alpha \eta(t) }{T \gamma(t) } (1 - M \eta(t)(1 - \alpha) ) || \exp(-\tau(t) ) ||^2$ 

$ + \sum\_{t=0}^{T-1} \frac{M^2 \eta(t)^2n\_{in}^{max} }{2 T \gamma(t)} G^2 $ + $\sum_{t=0}^{T-1}  \sum_{k=1}^N \frac{\mathbb ||\tau(t) - \tau\_k(t)||^2 }{NT\gamma(t)}$.

In SpaFL, we enforce sparsity through the sparsity regularizer $\alpha$. We can see that as $\alpha$ increases, the convergence rate can be damaged due to the second and the last terms on the right-hand side. Hence, a large $\alpha$ can induce more sparsity to models, but it can damage the performance. We can also see that the convergence rate depends on the difference of gradients between global and local thresholds as shown in the first term, thereby capturing the impact of the non-iid dataset. We will provide the detailed proof in the appendix of the revised version, 

[1] Li, Xiang, et al. "On the convergence of fedavg on non-iid data." arXiv preprint arXiv:1907.02189 (2019).

[2] Yi, Kai, et al. "FedP3: Federated Personalized and Privacy-friendly Network Pruning under Model Heterogeneity." arXiv preprint arXiv:2404.09816 (2024).
Thank you for your rebuttal and for clarifying the threshold settings and the generalization bound in your paper. However, I still have concerns regarding the theoretical analysis and comparison with baseline methods.

**Convergence Rate and Communication Complexity**: My concern pertains to the lack of theoretical analysis concerning the convergence rate and communication complexity of your proposed method. Currently, your paper and theorem do not include an analysis or theoretical comparison with existing works. I am looking for a more comprehensive analysis and a detailed theoretical comparison with established studies.

**Baseline Comparison**: While I appreciate the inclusion of baselines such as FedSpa (2022) and FedP3 (ICLR24) in your rebuttal, the main comparison with Fjord (NeurIPS21) and HeteroFL (ICLR21) in your paper and FedSpa (2022) still seems outdated for a NeurIPS24 submission. Regarding FedP3, I understand the rationale behind selecting a global pruning rate of 0.5. However, I must point out a discrepancy. FedP3 employs personalized model aggregation not for enhancing performance but for significantly reducing communication costs and maintaining privacy. Comparing this with your method, which uses full-layer aggregation, to FedP3’s partial-aggregation is *potentially unfair and misleading*. It is essential to clarify this point. Since the FedP3 study demonstrated that reducing the global pruning rate from 0.9 to 0.5 significantly impacted performance, the results you presented are not surprising. I would like to see a comparison that includes trends in communication costs between FedP3 and your method. Additionally, what would the performance of FedP3 look like with a higher global pruning rate? Alternatively, as suggested in my initial review, consider comparing your method with more recent baselines to clearly demonstrate the effectiveness of your method.
Thank you for your rebuttal and additional experiments. I raised my score to 6 since my questions are well addressed.
As the reviewer mentioned, the current algorithm can be applied to transformer models. Since multi-head attention modules consist of multiple linear matrices: query, key, and value, we can use the same pruning approach as we did with linear layers in SpaFL. For each column of query, key, value matrices, we can define a trainable threshold, If the average magnitude of parameters in a column is smaller than its threshold, we can prune the entire column. We tested this with a vision transformer with depth of 6, heads of 8, dim of 128, and linear layer dimension of 256 on non-iid CIFAR-10 dataset over 100 clients with dirichelet distribution of 0.1. We briefly used the same hyperparameters as we did in CIFAR-10 experiment in our manuscript without hyperparameter optimization and averaged the results over three random seeds.

| Algorithm   | Accuracy (CIFAR-10) | Density|
| -------- | ------- | --------|
| SpaFL  | $68.6 \pm 1.5$    | 54.7%|
| FedAvg | $59.2 \pm 0.4$     | 100%|

Hence, the proposed algorithm can be applied to transformer models.
I have read the rebuttal. I think the paper contains interesting insights on the difficult problem of sparse training so I am keeping my score.