Assortment of Attention Heads: Accelerating Federated PEFT with Head Pruning and Strategic Client Selection
Authors: Yeshwanth Venkatesha
Souvik Kundu
Priyadarshini Panda
Venue: Accepted by TMLR ()

Abstract: Parameter Efficient Fine-Tuning (PEFT) has become the de facto approach in adapting Large Language Models (LLMs) for downstream tasks in Natural Language Processing. However, its adoption in privacy-preserving distributed learning frameworks, such as Federated Learning (FL), remains relatively limited. This is mainly due to challenges specific to FL, such as resource-constrained devices and diverse data distributions among clients. In this paper, we propose an efficient method to perform PEFT within the FL framework for Multi-Head Attention (MHA) based language models. We address the challenges through head pruning, a novel head-specific weighted aggregation mechanism, and a client selection strategy. Head pruning minimizes training complexity within the clients, guided by the importance score computed based on the confidence of the attention head. Weighted aggregation of heads ensures the global model captures crucial updates from diverse clients, complementing our client selection strategy. We show results on the MultiNLI benchmark along with 20 Newsgroups, XL-Sum, and E2E NLG datasets. We use the MultiNLI dataset and T5-small model with LoRA as our PEFT method, attaining sparsity levels of up to 90\%, resulting in a communication advantage of up to 1.8x and a reduction in training OPs of 3.9x while maintaining the accuracy drop under 2\%.

Published in Transactions on Machine Learning Research (10/2025)
Assortment of Attention Heads: Accelerating Federated PEFT with Head Pruning and Strategic Client Selection
Yeshwanth Venkatesha
Deportment of Electrical Engineering Yale Universily
Souvik Kundu
souvikk Gintel. com
Intel Labs
Priyadarshini Panda
priya.panda@yyle..com edu
Department of Electricol Engineering Yale University
Reviewed on OpenReview:
Abstract
Parameter Efficient Fine- Tuning (PEFT) has become the de-facto approach in adapting Large Language Models (LLMs)) for downstream tasks in Natural Language Processing. However, its adoption in privacy-preserving distributed learning frameworks, such as Federated Learning (FL), remains relatively limited. This is mainly due to challenges specific to FL, such as resource-constrained devices and diverse data distributions among clients. In this paper, we propose an efficient method perform PEFT within the FI. framework for Multi- Head Attention (MHA) based language models. We address the challenges through head pruning, novel head-specific weighted aggregation mechanism, and client selection strategy. Head pruning minimizes training complexity within the clients, guided by the importance score computed based on the confidence of the attention head. Weighted aggregation of heads ensures the global model caplures crucial updates from diverse clients complementing our client selection strategy. We show results on the MultiNLI benchmark along with 20 Newsgroups, XL-Sum, and E2E NLG datascts. We use the MultiNLI datasct and T5-small model with LoRA as our PEFT method, attaining sparsity levels of up 90%, resulting in a communication advantage of up to 1.8x and reduction in training OPs of 3.9x while maintaining the accuracy drop under 2%.
1 Introduction
[Brown et al. 2020). T5 (Raffel et al. 2020), BART Lewis which are built on the Multi-Head Attention (MHA) based transformer architecture Vaswani et al. 2017 have achieved exceptional outcomes a wide range of Natural Language Processing (NTP) tasks. Morcover, they have begun venture into other domains, including Computer Vision (CV) with models like VIT (Dosovitskiy et al.) 2020), Stable Diffusion (Rombach al. )022), and LayoutLM (Xu et al. 2020), well as Audio with models like Whisper (Radford et al. 2023 and XLS-R (Babu et al. 2021). Fine-tuning these models for downstream tasks such text classification, summarization, and machine translation has proven to be highly effective. Specifically, class of methods known as Parameter Efficient Fine-Tuning (PEFT) focuses on selectively updating a small subset of model parameters while keeping most of the pre-trained model unchanged (Hu et al.) 2021 Li & Liang 2021 Liu et al. (2021b). This approach significantly lowers the computational resources required language models (LLMs) to specific downstream tasks. However, as the data sources become edge-focused and individuals grow more privacy-conscious, there is shift. towards privacy-presee preserving distributed learning
Large Language models (LLMs), such as GPT
et al. 2019) and BERT Devlin et al. 2018
adapt large

---

Published in Transactions on Machine Learning Research (10/2025)
paradigms, such Federated Learning (FL) (McMahan et al.
2017 for several applications (Hard al.
2018 Chen et al. 2019: Liu et al.) 2021a: Babakniya et al. 2023b). Naturally would need to do federated
learning for language model fine tuning.
Challenges
Reduce training complexity
Clert
Customized training each client
Server
Network Challenge
Business Systems with
Voice Assistants with conversational
Reduce communication complexity
formal data
Servenside Chillenges
Ciient
P
Efficient Aggregation
40
Strategic Client Selection
Resource Constraints
E-readers with trending news articles
Smartphones with unique texting patterns.
P
Processor
Storage
D
Battery
Network BW
Figure 1: Within federated environment, devices often face resource limitations and exhibit variability in both data and available resources. Numerous challenges arise al each stage when finetuming a language model under these conditions.
As illustrated in Fig. ! consider scenario of FL involving distributed sources of text data. Each source exhibits unique text patterns, encompassing diverse text styles. For instance, business systems may contain formal documents, while voice assistants may comprise casual conversational data. Furthermore, each client is constrained by individual device limitations, such as battery life, processing power, storage, and connectivity. The process of training model in such an environment poses challenges on multiple fronts. On the client. side, constraints arise duc to limited and often diverse resources across clients. On the server side, efficient aggregation and client selection stralegies become imperalive to handle client heterogeneily and accelerate convergence. Additionally, the communication cost between the client and the server constitutes an additional overhead.
Transformer models with MHA are shown to effectively capture intricate textual details, as specific heads encode syntax, rare words, and positional nuances. Additionally, it has been demonstrated that several heads are redundant and can be pruned without sacrificing the accuracy (Voita et al., 2019 Li et al. 2021 Behnke & Heafield 2020 Shim et al. 2021). In this work, we investigate harnessing the inherent structure offered by multiple heads in MHA mechanism to devise an efficient FT. PEFT method. Our goal is to minimize the training and communication complexity for individual clients while facilitating a cohesive aggregation of models independently trained on their private data. To this end, we first implement head pruning to reduce the training and communication complexity at the client level. The pruning decision is guided by the importance score specific to each client's heads. The importance score works as a fingerprint of the data present in the clients. Next, we perform weighted averaging of the retained heads, ensuring that the server accentuates the information from unique data distributions from all participating clients.
Further, in typical large-scale FL systems, the sheer number of devices makes it essential to select subset of devices for each FL round. This client selection directly impacts the performance of the system, making a tailored approach necessary. We propose a client selection strategy based on the difference in clients' local loss and global model loss prioritizing the updates from clients with maximal impact. The weighting mechanism of attention heads proves especially beneficial when limitod number of clients are chosen for training. Put simply, our approach addresses the training and management of multiple attention heads trained from distributed data sources, hence the title "Assortment of Attention Heads". Our contributions are summarized as follows.
To minimize training and communication complexity at the clients, we implement a head pruning strategy using head importance scores determined by the maximum attention scores generated by each head.
We introduce an aggregation mechanism in which updates from each head are weighted according their importance, allowing for the collection of the most crucial updates to the global model. This
2

---

Published in Transactions on Machine Learning Research (10/2025)
effectively mitigates the performance drop caused by sparsity, thereby enhancing the convergence rate.
We propose client selection strategy based the loss difference between the global model and the client model to further improve performance.
We demonstrate the efficacy of our approach on text classification datasets, MultiNLI and 20 Newsgroups, as well on more demanding summarization task using the XT.Sum dataset and generation task using E2E NLG dataset. On MultiNLI dataset, achieve 90% head pruming translating to 1.8 times and 3.9 times reduction in communicttin complexity and training operations (OPs) respectively compared to training the full network with standard FedAvg. Additionally, we achieve faster convergence compared to random client selection at similar levels of sparsity.
2 Background and Related Work
Muti-Hcad Attention: The standard attention mechanism in transformers VVawwnni et al. 2017 is defined
88:
QKT Vox
At, K, softmax
V
(1)
where Q K, and V correspond to Query, Key, and Value matrices respectively. In practice, the transformer model uses multiple attention heads to capture different aspects of the relationships between words. Let h denote the index of the attention head. Each attention head computes a different set of query, key, and value vectors. The outputs of all attention heads are then concatenated and lincarly transformed to obtain the final output. The multi-head attention (MHA) operation is defined as follows:
MHA K, V) = Concat(headh,... head,) Wo
(2)
where
head, = Attn(Q WQr: K- WKh, V Wvn)
(3)
Here, Wor,, WK,, WVn, and Wo are learnable weights corresponding to each attention head.
Parameter Efficient Fine-Tuning Transfer learning has been pivotal approach in leveraging pre-trained models for various natural language processing tasks. Rescarchers have explored techniques for fino-tuning pre-trained models specific downstream tasks to achieve improved performance with fewer computational resources. The umbrella term for this is Parameter Efficient Fine Tuning (PEFT). Adapters (Houlsby et al., 2019), LoRA (Hu et al. 2021), AdaLoRA (Zhang et al. 2023a). 1A3 (Liu et al. 2022a), BitFit (Zaken et al, 2021), and Ladder Side Tuning (LST) (Sung et al. 2022) represent some of the widely adopted techniques in this domain. Prefix tuning stralegies involve adjusting the model's prompl or prefix to tailor its behavior for specific tasks (Li & Liang, 2021 Liu et al.) 20211 2023 Lester al. 2021]. Our approach is not limited to a specific PEFT strategy; rather can be applied to any PEFT strategy long its trainable paramcters can be associated with the heads of the MHA mechanism.
Federated Learning with Transformers: Previous works have explored federated learning methods specifically catered for transformer models. There are several works that benchmark NLP tasks that primarily rely on transformer models within the federated learning environment, as evidenced by (Hilmkil et al.
2021
Lin et al.] 2021). In an alternative research direction, personalized federated learning approaches using transformer architectures have been explored for customizing models to individual clients (Li et al. 2023
Sun et al. 2023 Kim et al., 2023). Simultaneously, in computer vision, works have compared CNNs and
transformer models for vision tasks in federated learning (Qu et al.
2022
Chen et al. 2022 Zuo et al. 2022
Specifically addressing PEFT in foderated lcarning, (Hycon- Woo et al.
2021
Babakniya et al. 2023a Yan
el al. 2024 leverage the LoRA framework for efficiency, and (Zhao el al.
2022
sludies prompt tuning ellects.
Various studies benchmark diverse PEFT methods in federated learning. covering privacy and resource constraints (Zhuang et al.] 2023 Zhang et al. 2003bb [Sun al.) 2022). While prior works adapt transformers to the federated learning paradigm, they have not exploited the natural structure of the multi-head attention mechanism for training and communication efficiency.
3

---

Published in Transactions on Machine Learning Research (10/2025)
Heads
Heads
MIA Layers
MHS& Layers
params (9)
Trainable adapter params (()
Pruned adapter
Server
Head Importance
Weighted Aggregation based on head
importance score
2
2
Client Selection based on loss difference to the global model
Attention Head
Pruning based on client-specific head
2
2
8
2
2
2
2
importance score
2
2
Client
Lgg
Figure 2: An overview of our stralegy encompassing al attention head pruning guided by importance scores and weighted aggregation, alongside client selection determined by the difference in loss between the global server model and the local client model.
Pruning in Transformers: Pruning a well-known process in machine learning where the redundant parts of the model are removed to achieve efficient processing (Sietsmaa & Dow 1988
LeCun et al. 1989
Han et al. 2015 Blalock al.] 2020)- Although initially popularized in the vision domain, adaptations of pruning techniques have been applied to transformer models well (Mao et al.
2021
Lagunas et al. 2021
The concept of head pruning, as presented in the 'Story of heads' (Voita et al.
2019
demonstrates that a
significant number of heads can be pruned without sacrificing performance, a notion extended by subsequent works such as (Held Yang. 2022; Wang et al. 2021h. The authors of (Gordon et al. 2020) study the effect of pruning the backbone model on fine-tuning. We take inspiration from previous work and adapt to federated settings.
Client Selection in FL: Federated learning typically assumes unbiased client selection, randomly choosing subset for each round. However, biasing this selection can potentially improve convergence rates and mitigate non-lID data concerns (Cho et al.] 2022 Zhang et al., 2021 Fu et al. 2023). Some works adopt a systems approach, selecting clients based on resource availability, particularly in IoT applications (Nishio & Yonetani 2019 Xu & Wang 2020 AbdulRahman et al., 2020). The client selection algorithm in (Huang et al. 2020) incorporates fairness considerations. Prior approaches often rely on pre-communication local training, or resource monitoring. In contrast, our method selects clients based on the most recent loss value of their local models, prioritizing clients whose updates are expected to provide the highest impact the global model. This adaptive, loss-based client selection, when combined with allention-head pruning and seworrtanceweiggtee aggregation, constitutes novel system-level contribution that improves convergence efficiency while remaining computationally lightweight.
3 Methodology
Consider a PEFT sctup with model M (0, b) where ④ the frozen backbone parameters and & is the set of trainable parameters (See Fig. 2 A federated learning system consists of a server and clients indexed with (C], where each client e contains a local dataset D.. The federated learning process spans over R rounds of communication. In round r, the server selects subset of clients S(r) (<< C), and broadcasts the trainable parameters ((r-1) for local training. The objective of each client is to minimize the local loss La while updating only the local trainable parameters h.::
arg min L. 4.).
(4)
After the clients complete their local training over a certain number of epochs, the server gathers their model updates and integrates them into the server model as follows:
()) = ((r-1) 1) + E EcEgoo) IDel (r.
(5)
Ecesin) Del
4

---

Published in Transactions on Machine Learning Research (10/2025)
Here, A Φ6(") (((-1) is the model update from client at round r, and 1/ is the server learning rate controlling global model changes each round. The model updates are weighted by their respective data sample count.
We identify the inefficiencies in this framework by asking the following key questions:
1. Is it necessary to train all the parameters across all clients? Can we achieve comparable performance by selectively training a subset of Φ_?
How can we modify the aggregation mechanism with custom weight to each parameter to maximize
the collective knowledge aggregated from different clients?
3. What strategy can we employ to optimize the selection of s for maximizing the performance?
In the following sections, we lackle these concerns through our proposed stralegies. As illustrated Fig. we tackle the problem with three key methodologies. First, we reduce the training and communication demands of each client by performing head pruning which we describe in Section 3.1 Next, we propose novel weighted aggregation technique to maximize the knowledge transfer from pruned models (refer to
Section 3.2 Finally, we introduce client selection approach to improve the convergence rate (refer to Section
3.3
3.1 Head Pruning
Assuming multi-head attention model with H heads, designate a subset of trainable paramcters P C + such that P = Pp...P.....P PH. h: each associated with a head indexed by [ht]. We prune away a subset of these parameters to increase the efficiency of the local training as well as the communication cost (see Fig. 2 To determine which parameters to prune, we compute the importance of each attention head at a given client by averaging the maximum attention scores (excluding the end-of-sentence symbol) across the client's local dataset (Voita et al., 2019). This score reflects the "confidence" of an attention head in its predictions, which proxy for its importance. The resulting importance score also serves fingerprint of the client's data distribution. Formally, the importance score of head h at client is defined as:
Σ max (()) (t). ahe Deel IEDe
(6)
where dh (1) and (1) max is taken over the token positions in the sequence, and the final score is averaged across all samples in
(1) denote the query and key vectors of head h for token in the i-th data sample. The
the client dataset D... We then prune the heads whose importance is below certain threshold calculated according to the desired level of sparsity as:
ahe if ane threshold otherwise
ahe
Note that the values of P corresponding to the pruned values of a need not be communicated to the server thereby saving a significant amount of communication cost.
We opt for Low-Rank Adapters (LoRA) (Au et al., 2021) as our chosen PEFT approach. In LoRA, the gradient of cach weight matrix in the MHA layer decomposed into two low-rank matrices, denoted as AW BA. Our updates then applied exclusively to these parameters and the designated task head, denoted T. Consequently,
trainable parameter set is denoted as & B.T).
For an input token x with dimensionality of d, the matrix A maps it to lower dimension d', and this lower-dimensional representation is then mapped back to the original dimensionality of by matrix B. Since multi-head attention operates along the d dimension, we can distribute only the matrix B across the attention heads. Thus, in this context, the set P p= (B) Note that while we illustrate the application with LoRA, our approach can be extended to any PEFT method where trainable parameters can be mapped specific attention heads.
5

---

Published in Transactions on Machine Learning Research (10/2025)
3.2 Weighted Aggregation
At round r, let the model update corresponding to head h from each client be denoted by Aphe (r)
(r) We weight
(r) Appe by their importance score the to update the corresponding parameters in the global model Pig follows:
(r)
(r) Phg = Phg (r-1) + Ecegin (r)
the Ahh
Eeeson (r)
Ohe +
is small constant added for numerical stability and 1j is the server learning rate controlling the rate of global model updates. The weighting of updates based on the importance of each head for a specific client ensures that the global model extracts the maximum benefit from the individual client contributions. The remaining trainable parameters that cannol be assigned to an MHA head (the set $ P) are updated following Equation 5
3.3 Client Selection
To determine participating clients S(r), we select the top K clients with the highest difference in loss compared to the global model. Using the server loss c(r-1) (the loss of the global model at the server) and each client's loss from their last communication L ((r)) calculate the difference and choose clients with the maximum disparity:
5(r) = arg max (C --6--1) 1.....K
The rationale behind this selection criterion that as this difference increases, it indicates that the client is falling behind the global model. Consequently, training updates from such clients expected to contribute significant improvements when incorporated into the global model. We initiate all clients with high loss values to promote initial engagement from clients that have not participated previously. We summarize the overall flow of our method in algorithm
Algorithm 1: Federated PEFT Algorithm
Input: Model M = {0, }). C clients, R rounds, learning rate
for r r= to R do
Select the set of participating clients ((r) (Eqn. 9
Broadcast ((r-1) to g(r);
for each client e S(r) do
Perform local training (Eqn. 4
Compute model updates; Compute importance score az; Communcaates sparse parameters P. and head importance scores
Server updates global model (Eqn.
4 Experiments
We consider the popular text classification datascts MultiNLI WWilliams et al., 2018 and 20 Newsgroup [Lang] 1995) and two generation datasets, XL-Sum (Hasan et al., 2021) and E2E NLG (Novikova al., 2017). ) MultiNLI dataset contains over 400k sentence pairs mapped to 3 classes entailment, contradiction, and neutral. 20 Newsgroup is a collection of nearly 20k documents spanning 20 different newsgroups ranging from technology and science to politics and sports. We take the English subset of XL-Sum dataset which contains over 330k samples of annotated article-summary pairs. E2E NLG is data-to-text generation task with 42k samples.
We uniformly distribute the training data among clients for each dataset and conduct evaluations on the held-out validation set.. We adopt publicly available splits for training and validation data in XL-Sum and 20

---

Published in Transactions on Machine Learning Research (10/2025)
Method
CS
Acc.
Conv Rate
Sparsity
OPs
Comm. Complexity
Rand
74.03 + 0.25 %
B1. FodAvg
0%
33.07G
3.38MB
Loss
78.76 + 0.45%
B2. FedArg +
Rand
71.39 = 0.45 %
90%
8.47G
1.86MB
Rand Pruning
Loss
76.87
0.39 %
Rand
72.76 + 0.38%
9 90%
Ours
8.47G
1.86MB
Loss
76.94 + 0.34 %
11
Table 1: Ablation study on the client selection (CS) and head pruning strategics on MultiNLI dataset with 100/2 clients. For each of the baselines we show results for both random client selection (Rand) and strategic losf-difference based clients selection (Loss). In addition to accuracy, also provide information on the convergence rate, quantified by the number of rounds needed to reach a. 70% accuracy threshold.
Newsgroup. In the case of MultiNLI, we merge both matched and mismatched validation sets into a unified validation set. Across all experiments, a batch size of 80 and local learning rate of 50-4 are utilizod. We employ local epoch and conduct 100 communication rounds. All the experiments are performed machine with one A100 GPU with 80GB GPU memory and an 8-core CPU with 16GB memory per core.
In addition to accuracy, we measure the computation and communication complexity follows:
Computation Complexity (OPs): Specifically, we account for the forward and backward compu- tation of Q, K, V matrices, attention layers, and the subsequent fully connected layers. We use the maximum token length of T 512 across our experiments. For simplicity, we omit the peripheral layers such as tokenizer, positional encoding, and final dassification/lagguaa model head since they common across all the methods. Details provided in A Communication Complexity: For communication complexity, we calculate the total number of Irainable paramelers in each client since only these needs to be communicated to the server each round.
We employ the MultiNLI dataset for a comprchensive analysis, evaluating the influcnco of our client selection approach and the significance of head pruning based on importance scores. We assess their effects on both the accuracy and convergence rate of the model, measured by the number of rounds required to achieve 70% accuracy. We establish the following baselines to compare the effectiveness of our design.
B1: Standard federated averaging without head pruming (FedAvg).
B2: Federated averaging with randomly selected heads for pruning (FedAvg + Random Head Pruning).
In our method we employ a head pruming of 90% based on the head importance score (Eqn. 6 the methods B1, B2 and ours we show results for both random client selection and loss based strategic client selection in Table We summarize the observations as follows:
For each of
Loss difference-based client selection directly influences the accuracy by approximately 5%, regardless of the chosen head pruning strategy, shown in the comparison between two consecutive rows.
indicated by the consistent Client selection does not impact the convergence rate independently,
mumber of rounds required to reach 70% accuracy in both B1 (5 rounds) and B2 (17 rounds).
Head pruning plays a significant role: transitioning from no sparsity in B1 to 90% sparsity with random pruning in B2 results in a noticeable decrease in both accuracy (~2%) and results in slower convergence (11 rounds).
Importance score based head pruning strategy alone enhances convergence rate (17 rounds in B2 9 rounds in ours) and improves accuracy by >1% compared to random pruning (B2).

---

Published in Transactions on Machine Learning Research (10/2025)
Dataset
# Clients
FedAvg (B1 + Random CS)
Ours
Comm.
OPs
Acc.
Comm.
OPs
Acc.
10/2
3.38MB
33.07G
78.34 + 0.23 %
1.86MB
8.47G
79.13 0.43%
MulGNLI
100/2
3.38MB
33.07G
74.03 + 0.25%
1.86MB
8.47G
76.94 0.34%
1000/20
3.38MB
33.07G
67.35+ 0.63%
1.86MB
8.47G
67.05 = 0.66%
10/2
5.06MB
228.90G
74.53 0.43%
2.78MB
78.34G
77.10 + 0.68%
Newsgroup
100/2
5.06MB
228.90G
72.01 + 0.56%
2.78MB
78.34G
77.59= 0.63%
1000/20
5.06MB
228.90G
67.89 1.03%
2.78MB
78.34G
75.56 0.83%
Dataset
# Clients
FedAvg (B1 + Random CS)
Ours
Comm
OPs
R1/R2/RI.
Comm.
OPs
R1 R1/R2/RT.
10/2
3.38MB
33.07G
234/0.12/0..3
1.86MB
8.47G
0.35/0.12/0.3
XL Sum
100/2
3.38MB
33.07G
0.33/0.12/0.33
1.86MB
8.47G
0.34/0.12/0.33
1000/20
3.38MB
33.07G
0.33/0.12/0.31
1.86MB
8.47G
0.34/0.12/0.31
10/2
3.38MB
33.07G
256/00.9/0.50
1.86MB
8.47G
0.55/0.32/0.51
E2E NLG
100/2
3.38MB
33.07G
0.55/0.29/0.50
1.86MB
8.47G
0.51/0.26/0.42
1000/20
3.38MB
33.07G
0.48/0.13/0.41
1.86MB
8.47G
0.47/0.13/0.42
Table 2: Results of our method across four datasets over different number of clients. For instance, in the configuration '100/2", a total of 100 clients exist, with 2 participating in each training round. We provide accuracy metrics for classification tasks and ROUGE metrics generation tasks. We show the baseline of using standard FedAvg without any sparsity and random client selection and compare il to our method with 90% sparsity and loss-based client selection.
for the
Combining head importance-based pruning with loss-based client selection yields higher accuracy and convergence rate compared to FedAvg with random pruning (B2) and achieves similar accuracy to FedAvg without pruning (B1).
We achieve 3.9-fold reduction computation complexity, as measured by OPs, and a 1.8-fold cnhancement in communication complexity compared to dense bascline (B1).
To delve deeper, we plot the validation accuracy during the initial 30 rounds of training in Fig. 3 The baseline B1 (FedAvg) with loss based client selection is the best performing among our experiments with higher accuracy as well as faster convergence. On the other extreme, baselinc B2 (FodAvg + random head pruning) with random client selection serves as the lower limit. We see that importance-based head pruning and weighted aggregation significantly improves the convergence rate, illustrating that our strategy enables positive collaboration among clients in the carly stages of training. However, as training progresses to the later stages where targeted model updates are crucial, the client selection strategy becomes more pivotal. Thus, both strategies complement each other, enabling training approaches that achieve performance levels close to the best-performing baseline scenario of no sparsity (B1) with loss-based strategic client selection.
We showcase the effectiveness of our methodology across four datasets, while considering diverse client configurations. We present accuracy as the primary performance measure for classification tasks, and ROUGE (OOUGE-/RROUGE-//ROUCCE
scores for summarization and summarization/ geneorationnnnnnnnnnooooooooooii tasks. To establish baseline for comparison, report the accuracy of standard FedAvg (B1 + Random CS). Our analysis encompasses scenarios involving 10, 100, and 1000 clients, each with a 2% participation rate, denoted as 100/2 and 1000/20 configurations. In the case of 10 clients, we consider the minimal participation of 2 clients. Across MultiNLI, XL-Sum, and E2E NLG lasks, we ut utilize T5-small as the backbone, while employing BART- for the 20 Newsgroup task. LoRA serves as the PEFT method across all datasets. We apply a head pruning of 90% according the calculated head importance scores. The aggregation process utilizes these scores and pruned gradients update the global model. Our method exhibits effectiveness, particularly in scenarios with high number of clients. Importantly, its applicability extends beyond conventional text.
8

---

Published in Transactions on Machine Learning Research (10/2025)
Training
FedAvg (B1
Random CS)
Ours
Method
Time
Commn.
OPs
Acc.
Max Pr%
Comm.
OPs
Acc.
FFT
47s
41.96MB
41.07G
75.60 0.63%
90%
23.08MB
10.63G
78.24= 0.59%
LoRA (r- 16)
38s
3.38MB
33.07G
74.03 + 0.25 %
90%
1886MB
8.47G
76.91= 0.34%
LoRA (r 8)
38s
1.69BB
32.38G
75.30:1 0.24%
90%
0.93MB
8.00G
76.67= 0.29%
LoRA (r 4)
38s
0.84MB
32.05G
75.22 0.45%
90%
0.46MB
7.76G
76.07= 0.43%
IA3
36s
168KB
38.70G
59.60 1.12%
50%
84KB
21.76G
60.14= 1.03%
Prompt Tuning
32s
743KB
22.81G
48.46 0.43%
90%
74KB
6.47G
53.93= 0.43%
P-Tuning
32s
1.80MB
23.80G
57.28 0.85%
90%
0.18MB
6.73G
60.83= 0.59%
Table 3: Effect of PEFT method. We compare several PEFT methods along with the full fine-tuning (FFT) baseline in the standard MultiNLI setting with 100/2 clients. We report training time per epoch, maximum sparsity before accuracy drops to random guessing levels (Max Pr%),, and communication cost (Comm.), MAC operations (OPs), alongside accuracy measurements (Ace.).
classification tasks, proving its effectiveness even in challenging tasks such as summarization and natural language generation.
Effect of Sparsity: Here, we investigate the impact of head pruming the performance of our method. In Fig. 4 we show the accuracy, number of OPs, and communication cost measured by the number of trained parameters sont by each client on an average per round. We vary sparsity across a range of sparsity levels, from 0% to 95%. Accuracy experiences only a minimal decline up to 90% sparsity; however, beyond 95%, the accuracy diminishes to level equivalent to random guessing. At 90% sparsity, observe 1.8x improvement in communication cost and 3.9x lower MAC operations with less than 2% drop in accuracy.
FL Training Curve - First 30 rounds
Accuiaa
70
60
50
B1 Loss- coss-based C5
82 Rand
(Rand
Effect of PEFT Method: In this study, we analyze the impact of various fine-tuning methods on our proposed approach. Specifically, we examine the outcomes of full fine-tuning (FFT), where all parameters of the backbone model are fine-tuned alongside several established PEFT methods including IA3 (Liu al.) 2022a),, P-Tuning (Liu et al. 2022b). and Prompt Tuning (Lester et al. 2021 Additionally, we present results for different ranks in LoRA. The findings are summarized in Table 3 which includes accuracy and system metrics training time, communication cost, and the OPs count per local round. We report the maximum sparsity achieved with our method before the model shows diverging behaviour. We include the bascline of standard FedAvg (B1 + Random Client Selection) for comparison. LoRA improves over full fine-tuning (FFT) in terms of OPs by 20% and > 12x in terms of communication cost with minimal accuracy drop. In our method, decrcasing the rank within LoRA leads to a steady decline in accuracy. While the improvement in computation complexity is marginal there is a proportional enhancement in communication complexity. This parameter can be tuned based on the specific application requirements. With IA3, Prompt tuning, and P- P-Tuning while we can achieve a significant reduction in communication cost, they are not competitive in the final accuracy compared to LoRA. Between P-tuning and prompt tuning, P-tuning demonstrates better performance owing to trainable tokens at every layer. However, it comes with the trade-off of increased communication complexity. Additionally, while IA3 is more efficient than LoRA with 440x less communication cost, it is also extremely sensitive to sparsity shown by the fact that we can only go up to 50%. Note that the training time can be further improved with a custom implementation of pruned model training.
(Lods-baeed
30
131
17
Round
0
5
10
15
20
Figure 3: Training curves over the first 30 rounds illustrating the impact of importance score- based head pruning and weighted aggregation. Baselines B1 and B2 are shown along with our method without strategic client selection. Pruning and aggrega- tion accelerate early convergence, while client selection improves later-stage performance.
9

---

Published in Transactions on Machine Learning Research (10/2025)
Training
Aed (B1
Random CS)
Ours
Model Type
Model Name
Time
Comm.
OPs
Acc.
Commn
OPs
Acc.
EncDec
T5-Small
38s
3.38MB
33.07G
74.03+ 0.25%
188MBB
8.47G
76.94 0.34%
EncDee
T5-
166s
10.13MB
128.19G
82.524
1.19%
5.57MB
34.8G
79.03 0.69%
EncDe:
BART-
82s
5.06MB
209.58G
66.804:
0.53%
2.78MBB
59.01G
78.77 0.54%
Enc-only
DistiBBERT
33s
6.22MB
42.73G
74.00
0.58%
3.42MB
11.61G
75.11 1.03%
Enc-only
RoBERTa
64s
7.89MB
21.89G
77.681 0.68%
4.34MBB
5.63G
82.24 0.43%
Dec-only GPT-2 Small
10s
3.94MB
42.73G
0.67/0.42/0.64
2.16MB
11.6GG
0.69/0.44/0.6
Table 4: Effect of different backbone architecture. Besides our usual experimental configuration using T5-small, we incorporate larger model represented by T5-base. Moreover, for comparative analysis, we incorporate both the base variant the BART model. We include two encoder-only models in DistilBERT and RoBERTa in addition to decoder-only model GPT-2. Our comparisons encompass training time per epoch, communication cost, and the maximum sparsity achieved without encountering model divergence. We use the standard configuration of MultiNLI datasel with 100/2 clients and LoRA as our PEFT method.
Heads
1.0
ID
Lavers
Distance
0.8
Non NDD-II
la - 0.11
0.6
0.4
Attn. Importance
Non-lia
0.7
0.0
Round
(b)
Figure 5: Assossment of head importance scores under IID and Non-IID scttings. (a) Head importance visualization for encoder-decoder cross-attention across two clients per round, during the first 10 rounds of training. (b) Quantitative comparison using Fuclidean distance between the two clients' importance scores in each round over the first 10 rounds.
Effect of Backbone Model: We also measure the influence of sub- stituting the backbone model on the overall system metrics in Table g Alongside the standard T5-small encoder-decoder model utilized in our experiments, we incorporate a larger T5-base model and BART [Lewis)
100
3:38 MB
2.55
1.06 MB
Accuracy (%)
80
60
93.07
et al. 2019). Furthermore, we introduce two encoder-only models, DistiBERTT (Sanh al.] 2019) and RoBERTa (Liu et al.) (019)) for the classification task MultiNLI. Finally, we include decodcr-only model GPT-2 (Radford et al. 2019) to show the performance on E2E NLG task. We employ the case of 100/2 clients and measure the same parameters training time per epoch, maximum sparsity achieved before the model diverges, communication cost, OPs, and accuracy. We find our approach consistently outperforms the baseline across all motries. Employing larger model cortainly improves performance, but comes with trado-offs. The training time increases by more than 4x, and the communication cost and OPs count triple when transitioning from T5-small to base Among encoder only models, RoBERTa provides better accuracy at lower compute complexity in terms of OPs. However, the communication complexity is higher than that of T5-small.
40.
Accuracy
1.77 MB
20
Comm. Cost
7.1.
20
40
60
80
100
Pruning %
Figure 4: Effect of pruning accu- racy, communication cost and OPs.
Effect of Non-IID data: We examine the impact of non-IID data in this section. We regulate the Non-IIDness by adjusting the a parameter the Dirichlet distribution. Employing our standard experimental setup with MultiNLI, 100/2 clients, and LoRA as the PEFT method, we present the results in Table Notably, even under extreme non-IID data conditions (a = 0.1), our method exhibits robustness, as we observe performance degradation. In Fig. 5 we offer qualitative and itative analysis of head
10

---

Published in Transactions on Machine Learning Research (10/2025)
importance scores for IID vs Non-IID scenarios. These scores are computed for each selected client over the initial 10 rounds of training. While both clients demonstrate similar scores in the case of TID (top 2 rows), a noticeable divergence appears in the case of non-lID, as illustrated in Fig. 5a To quantify this difference,
we plot the Euclidean distance between the scores for both scenarios Fig. 5b
The distance between the
clients low and nearly constant the case of ITD whereas the case of non- IID the distance significantly higher and has larger variation across rounds.
5 Conclusion and Discussion
this paper, we present our method to improve the training and communication complexity of train- ing PEFT in an FL environment. We accomplish this by capitalizing on the inherent structure within the multi-head mechanism of transformer models, employing a combination of attention serention-hedd pruning and strategic client selection. The empirical results demonstrate the efficacy of our approach across var- ious datasets such as MultiNL,I, 20 Newsgroup, XT. Sum, and E2E NLG. We conduct additional exper- iments to validate the effectiveness of our approach under diverse scenarios, encompassing non-IlI data settings, different PEFT methods, and backbone models.
Dirichlet Param a
FedAvg
Ours
IID
74.03= 0.25%
76.94 0.34%
2.0
74.61= 0.32%
77.70 0.26%
0.5
74.41= 0.13%
76.91 0.29%
0.1
74.45= 0.22%
77.65 0.34%
Table 5: Effect of Non-IID data. We vary the parameter of Dirichlet distribution to increase the degree of non-IIDness of the data. The results shown are for the standard case of LoRA on MultiNLI with 100/2 clients.
While our exploration limited to head pruning, we acknowledge numerous other avenues for enhancing model efficiency, such as token pruming, individual weight pruning, quantization, linear attention, and more. The challenges associated with integrating these methods into a federated environment remain an ongoing and open problem that merits further exploration. Our research is confined to empirical analysis and does not provide formal convergence analysis. Conversely, we have shown results by simulating federated environment. Showcasing it in real-world applications on low-rcsource devices, such as mobile phones, would expedite implementation. The applicability of our method beyond language lasks, such as vision transformers is untested.
References
Sawsan AbdulRahman, Hanine Tout, Azzam Mourad, and Chamseddine Talhi. Fedmces: Multicriteria client selection model for optimal iot federated learning. IEEE Internet of Things Journal, 8(6):4723 4735, 2020.
Sara Babakniya, Ahmed Roushdy Elkordy, Yahya H Ezzeldin, Qingfeng Liu, Kee-Bong Song, Mostafa El-Khamy, and Salman Avestimehr. Slora: Federated parameter efficient fine-tuning of language models. Xixiv preprint ar orX::2308800:: 2023a.
Sara Babakniya, Souvik Kundu, Saurav Prakash, Yue Niu, and Salman Avestimehr. Revisiting sparsity hunting in fedcrated learning: Why does sparsity consensus matter? Transactions on Machine Learning Research, 2023b.
Arun Babu, Changhan Wang, Andros Tjandra, Kushal Lakhotia, Qiantong Xu, Naman Goyal, Kritika Singh, Patrick von Platen, Yatharth Saraf, Juan Pino, et al. XIs-r: Self-supervised cross-lingua speech representation learning at scale. arXiv preprint arXie:2111.09200 2021.
Maximiliana Behnke and Kenneth Heafield. Losing heads in the lottery: Pruming transformer attention in neural machine translation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 2664 2674, 2020.
Davis Blalock, Jose Javier Gonzalez Ortiz, Jonathan Frankle, and John Guttag. What is the state of neural network pruning? Proceedings of machine learning and systems, 2:129-146, 2020.
11

---

Published in Transactions on Machine Learning Research (10/2025)
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-190 2020.
Jinyu Chen, Wenchao Xu, Song Guo, Junxiao Wang, Jie Zhang, and Haozhao Wang. Fedtune deep dive into efficient lederated fine-tuning with pre-trained translormers. arxiv preprinl arXiv:2211.08025, 2022.
Mingaing Chen, Ananda Theertha Suresh, Rajiv Mathews, Adeline Wong, Cyril Allauzen, Françoise Beaulays, and Michael Riley. Federated learning of n-gram language models. arXiv preprint arXiv:1910.08./ 03432, 2019.
Yae Jee Cho, Jianyu Wang, and Gauri Joshi. Towards understanding biased client selection in federated learning. In International Conference on Artificial Intelligence and Statistics, pp. 10351 10375. PMLR, 2022.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition al scale. aix preprinl arXiv:2010. 119299 2020.
Lei Fu, Huanle Zhang, Ge Gao, Mi Zhang, and Xin Liu. Client selection in federated learning: Principles, challenges, and opportunities. IEEE Internet of Things Journal, 2023.
Mitchell A Gordon, Kevin Duh, and Nicholas Andrews. Compressing bert: Studying the effects of weight pruning on transfer learning, arXiv preprint ar arXi::000.08807 2020.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXie preprinl urXi::1510.00149, 2015.
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chlos Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint 1oXi::811.00000 2018.
Tahmid Hasan, Abhik Bhattacharjee, Md. Saiful Islam, Kazi Mubasshir, Yuan-Fang Li, Yong-Bin Kang, M. Sohel Rahman, and Rifat Shahriyar. XL-sum: Large-scale multilingual abstractive summarization for 44 languages. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pp. 4693-4703 Online, August 2021. Association for Computational Linguistics. URL https://aclantholooy org/2021. http/222..findnggs-aa findings-acl.47 413
William Held and Diyi Yang. Shapley head pruning: Identifying and removing interference in multilingual transformers. arXiv preprint arXim:2210.0511 2022.
Agrin Hilmkil, Sebastian Callh, Matteo Barbieri, Leon René Sütfeld, Edvin Listo Zec, and Olof Mogren. Scaling federated learning for fine-tuning of large language models. In International Conference on Applications of Natural Language to Information Systems, pp. 15-23. Springer, 2021.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficien transfer learning for nlp. In International Conference on Machine Learning, pp. 2790-2799, PMLR, 2019.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:216..00 09685, 2021.
Tiansheng Huang, Weiwei Lin, Wentai Wu, Ligang He, Koqin Li, and Albert Y Zomaya. An efficiency-boostee client selection scheme for federated learning with fairness guarantee. IEEE Transactions on Parallel and Distributed Systems, 32(7):1552-1555 1564, 2020.
12

---

Published in Transactions on Machine Learning Research (10/2025)
Nam Hyeon- Woo, Moon Ye-Bin, and Tao-Hyun Oh. Fedpara: Low-rank hadamard product for communication- efficient federated learning. arXiv preprint arXiv:21 008.06098, 2021.
Yeachan Kim, Junho Kim, Wing-Tam Mok, Jun-Hyung Park, and SangKeun Lee. Client-customized adaptation for parameter-effiiint federated learning. In Findings of the Association for Computational Linguisties: ACL 2023, pp. 1159 1172 2023.
François Lagunas, Ella Charlaix, Victor Sanh, and Alexander M Rush. Block pruning for faster Transformers. arXiv preprint arXit:2109.04880 2021.
Ken Lang. Newsweeder: Learning to filter netnews. In Armand Prieditis and Stuart Russell (eds.), Machine Learning Proceedings 1995, pp. 331 339. Morgan Kaufmann, San Francisco (CA), 1995. ISBN 978-1-55860- 377-6. doi: https://doo..org10.... 1tt.//1..11111111111 1tt.g//0...011///8--555555600000 1-778----5800377--6665
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-effiiient prompt tuning. arXiv preprint ar arxix:2104 08691, 2021.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising generation, translation, and comprehension. arXiv preprint arXi:: 19461, 2019.
pre-training for natural language
Hongxia Li, Zhongyi Cai, Jingya Wang, Jiangnan Tang, Weiping Ding, Chin-Teng Lin, and Ye Shi. Fedtp: Federated learning by transformer personalization. IEEE Transactions Neural Networks and Learning Systems, 2023.
Jiaoda Li, Ryan Cotterell, and Mrinmaya Sachan. Differentiable subset pruning of transformer heads. Transactions of the Association for Computational Linguistics, 9:1442-1459, 2021.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference
Natural Language Processing (Volume 1: Long Papers), pp. 4582-4597, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1 11222/1.c2lll////
URL https::/acclattholggg. hwtp:tthhooooooooog..ooooo www.or/gg........// 222//22211.ccc acl- 1ong.3333333 353
Bill Yuchen Lin, Chaoyang Hc, Zihang Zeng, Hulin Wang, Yufen Huang, Christophe Dupuy, Rahul Gupta, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. Fedulp: Benchmarking federated learning methods for natural language processing tasks. arXiv preprint arXiv:2104.0888 088155, 2021.
Haokun Liu, Derek Tam, Mohammed Muqecth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning, Advances in Neural Information Processing Systems, 25:1950-1965, 2022a.
Ming Liu, Stella Ho, Menggi Wang, Longxiang Gao, Yuan Jin, and He Zhang. Federated learning meets natural language processing: A survey. arXiv preprint arXiv:2107. 12603, 2021a.
Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. P-tuming v2: Prompt tuning can be comparable to finc-tuning universally across scales and tasks. orXiv preprint arXiv:2110.07660 2021b.
Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang, P-tuning: Prompt tuning can be comparable to fine-tuning across scales and tasks. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 61 68, 2022b.
Xiao Liu, Yanan Zheng, Zhengxiao Du, Ming Ding, Yujie Qian, Zhilin Yang, and Jie Tang. Gpt understands, too. AT Open, 2023.
13

---

Published in Transactions on Machine Learning Research (10/2025)
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint Ar::::1907. 2019.
Jiachen Mao, Huanrui Yang, Ang Li, Hai Lí, and Yiran Chen. Tprune: Flffccient transformer pruning for mobile devices. ACM Transactions Cyber-Physical Systems, 5(3):1 22, 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Soth Hampson, and Blaise Aguera y Arcas. Communication- efficient. learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp. 1273 1282. PMLR, 2017.
Takayuki Nishio and Ryo Yonetani. Client selection for federated learning with heterogeneous resources in mobile edge. In ICC 2019-2019 IEEE international conference on communications (ICC), pp. 1-7. IEEE, 2019.
Jekaterina Novikova, Ondiej Dušek, and Verena Rieser. The e2e dataset: New challenges for end-to-end generation. arXiv preprint arXiv: 09254, 2017.
Liangqiong Qu, Yuyin Zhou, Paul Pu Liang, Yingda Xia, Feifei Wang, Ehsan Adeli, Li Fei-Fei, and Daniel Rubin. Rothinking architecture design for tackling data heterogencity in federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pullern Recognition, pp. 10061 10061-10071, 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodci, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl I blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision. In International Conference Machine Learning, pp. 28492 28518. PMLR, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Reseurch, (11):5485--555 2020.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910. 110.01100, 2019.
Kyuhong Shim, Iksoo Choi, Wonyong Sung, and Jungwook Choi. Layer-wise pruning of transformer attention heads for efficient language modeling. In 2021 18th International SoC Design Conference (I80CC), pp- 357-358. IEEE, 2021.
Sietsma and Dow. Neural net pruning-why and how. In IEEE 1988 international conference on newral networks, pp. 325 333. IEEE, 1988.
225-333.
Guangyu Sun, Matias Mendicta, Taojiannan Yang, and Chen Chen. Exploring parameter-efficient fine-tuning for improving communication efficiency in federated learning. arXiv preprint arXiv:2210.01 708, 2022.
Guangyu Sun, Matias Mendieta, Jun Luo, Shandong Wu, and Chen Chen. Fedperfix: Towards partial model personalization of vision transformers in federated learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4988-4998, 2023.
Yi-Lin Sung, Jaemin Cho, and Mohit Bansal. Lst: Ladder side-tuning for parameter and memory efficient: transfer learning. Advances in Newral Information Processing Systems, 35:12991-13005 2022.
Ashish Vaswani, Noam Shazoer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
14

---

Published in Transactions on Machine Learning Research (10/2025)
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self- attention: Specialized heads do the heavy lifting, the rest can be pruned. In Anna Korhonen, David Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics. doi:
10.18653/v1. / 1580. URL
Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade loken and head pruning. In 2021 IEEE International Symposium on HiihhPPrrformrce Computer Architecture (HPCA), pp. 97-110. IEEE, 2021.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapler of the Association for Computational Linguisties: Human Language Technologies, Volume 1 (Long Papers). pp. 1112-1122. Association for Computational Linguistics, 2018. URL http://aclweb or org/ //aggtoooggy////// N18-1101
Jie Xu and Heqiang Wang. Client selection and bandwidth allocation in wireless federated learning networks: A long-term perspective. IEEE Transactions on Wireless Communications 20(2):1188-122 2020.
Yiheng Xu, Minghao Li, Lei Cui, Shaohan Huang, Furu Wei, and Ming Zhou. Layoutlm: Pre-training of text and layout for document image understanding. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp.
2020.
Yuxuan Yan, Qianqian Yang, Shunpu Tang, and Zhiguo Shi. Federa: Efficient fine-tuning of language models in federated learning leveraging weight decomposition. arXiv preprint arXie:2404.18844 2024.
Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-effiiient fine-tuning for transformer-basee masked language-models. arXiv preprint arXiv:2106. 101999, 2021.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficien fine-tuning. arXiv preprint arXiv:2308.1051 2023a.
Wenyu Zhang, Xiumin Wang, Pan Zhou, Weiwei Wu, and Xinglin Zhang. Client selection for foderated learning with non-iid data in mobile edge computing. IEEE Access, 9:24462 24474, 2021.
Zhuo Zhang, Yuanhang Yang, Yong Dai, Qifan Wang, Yue Yu, Lizhen Qu, and Zenglin Xu. Fedpetuning: When federated learning meets the parameter--ffiiient luning methods of pre-trained language models. In Annual Meeting of the Association of Computational Linguistics 2023, pp. 9963-99777 Association for Computational Linguistics (ACL), 2023b.
Haodong Zhao, Wei Du, Fangqi Li, Poixuan Li, and Gongshen Liu. Reduce communcctiin costs and preserve privacy: Prompt luning method in federated learning, arXie preprint arxiv:2008. Nix::2208.1268, 2022.
Weiming Zhuang, Chen Chen, and Lingjuan Lyu. When foundation model meets federated learning: Motivations, challenges, and future directions. arXiv preprint arXiv:2306.15544 2023.
Xiaojiang Zuo, Qinglong Zhang, and Rui Han. An empirical analysis of vision transformer and cmn in resource-cnnstrained federated learning. In Proceedings of the 2022 5th International Conference Machine Learning and Machine Intelligence, pp. 8 13, 2022.
A MAC Operations
The layerwise operations and their corresponding OPs count are detailed in Table 6 The main components are computing query (Q), keys (K) and values (V), attention layers followed by fully connected layers. For each of the PEFT methods, we provide the components specific to them. In Table 7 we provide a summary of the number of layers (L), attention heads per layer (H), and the hidden dimension of the model (d) for all considered models.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Layer
Fwd/Bwdd
Computation
OPs
Pwd
Y- WX
QKV
Bwd
4X dYW", .aw XYX
LIT TT)
Y -VVA== SMIQK*)
LIT T
Attention
Bwd
aV dYA. - dYV, dAK, dk dAQ
LIT9 7'd T'd
Y=WX WX
FC
Bwd
4X dVWT,at -dx
1TT8 (T)
Y= BZ, =
LTT8+ Tr
LoRA
Bwd
aZ B, dYZ,dA dZX,ex AZA
LTdr TTd Ter Trd+
Fwd
V-1,x
14Td) 4TT+Td)
IA3
Bwd
XdY, =
Prompl Tuning
Pwd
dx, dx,=d ,dV =dYA dV dAK,
Y, = W.Y. =VV,A=
T, + (T, 27 + (T, +1
Te + +2) (1,+ -T
Pod
v, - W.Y... V- -VV,A SMIQKT)
LT, (T, TYP (T, 119d)
P-Timing
=x, ad,WW,TTT A, dVV dAK dAQ
10758 (7, TPP +7, 77,47744
Table Layerwise computations and their corresponding OPs. X: Input, Y: Output, d: hidden dimension, T: Number of tokens, W: Weights, Z: Intermediate output, A: Attention output, A, B: Low-rank approximation of W, 4; Learnable scaling factors in IA3, Ty: Number of prompt lokens, Xp: Prompt token inputs, Yp: Prompt token outputs, L: Number of layers.
Model
L
H
T5-Small
18
8
512
T5-Base
36
12
768
BART
36
16
1024
DistilBERT
12
12
768
RoBERTa
12
12
512
GPT-2 Small
12
12
768
Table 7: Model configurations: Number of layers (T.), number of attention heads per layer (H), and hidden dimension (d).
16


Reviews:
[{"review_id": "J2wfvATVfg", "reviewer": "Reviewer_K38C", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper focuses on federated learning (FL), proposing to combine parameter-efficient finetuning (PEFT), weight pruning for multi-head attention (MHA) layers, and a loss-based client selection method in order to improve on the amount of local compute and communication required in training a suitable model. The authors experimentally demonstrate that their proposed approach reduces the amount of bits communicated as well as required training operations.", "claims_and_evidence_yn": "No", "claims_explanation": "I have several issues with claims and evidence in this paper:\n\n1) Some of the claimed contributions are essentially already known in the existing literature (especially loss-based client selection, see Requested changes for details).\n2) There are some important details missing from the experimental setup required for properly evaluating the work (see Requested changes for details).\n3) None of the experimental results have any error metrics, which makes it impossible to ascertain how reliable the results are w.r.t. inherent randomness in training.\n4) Some claims seem somewhat misleading (see Requested changes for details)\n5) The proposed approach, especially the client selection, does not seem to sit very well with the general setting as it is introduced in the paper (standard cross-device): as it is quite unlikely that training runs long-enough to go through all devices at least once, the loss-based client selection is basically just picking some unseen devices for (almost) the entire training.", "audience_yn": "Yes", "audience_explanation": "Yes. In general, improving the efficiency of FL is an important and interesting topic. The paper has, eg, nice ablation studies on how the different components affect the overall performance.", "requested_changes": "In decreasing order of importance:\n\ni) On the client-selection based on local loss: Cho et al. 2022: Towards Understanding Biased Client Selection in Federated Learning have already introduced and formally analyzed such a local loss-based client-selection method. Please clarify what exactly is your contribution here.\n\nii) How are all the hyperparameters tuned? Are these tuned separately for each method, including the baselines?\n\niii) Add some repeats to the experiments (if not done), and report some error metrics as well, eg, standard error of the mean (at the very least for some key experiments if this is not feasible for everything).\n\niv) Sec 4: How much compute do the peripheral layers omitted from the results require?\n\nv) Considering the claim of accelerating FL (eg, the title of the paper), this seems somewhat misleading, as based on Table 1, pruning is making learning a lot slower, while this is somewhat mitigated by using the proposed pruning instead of pruning randomly. Looking therefore, eg, at the total communication complexity (over rounds until convergence), the proposed method is pretty much on par with standard FedAvg without pruning. Can you comment on this?\n\nvi) Looking at the total number of clients in the experiments, the size of the client subset per round, and the total number of rounds, it seems like the loss-based client selection could effectively be just a round robin selection (eg with Table 1 experiments, first 50 rounds guarantee always selecting an unseen pair of clients, while for most of the last 50 rounds, would guess that the clients encountered early on during the first phase have higher loss than more recent ones). Can you confirm if this is true? Have you checked if the loss-based selection brings any benefit compared to round robin.\n\nvii) On head pruning: Voita et al. 2019 claim that the head confidence used in the proposed method is not a good way to do pruning, and instead propose Layerwise Relevance Propagation (LRP). Have you tested this or why do you still use the confidence?\n\nviii) Non-IID data, p10: please clarify how the client split is done, is this just creating clients with different number of samples (but otherwise IID data)?\n\nix) Server loss $\\mathcal{L}_g$ in Eq.(9) is not defined anywhere(?), please fix.\n\n\n### Minor tweaks (no comments needed about addressing these):\n\n* Please use citep/citealt/citet etc. to fit the citations to the text, making the paper smoother to read.\n\n* Please fix typos: \n    * p7 summary of observations: Head pruning with B2 should have convergence at 11, not 12 rounds to match Table 1?", "broader_impact_concerns": "No broader concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XlU7ecNCjn", "reviewer": "Reviewer_DxFc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents an approach for handling the training and management of multiple attention heads trained from distributed data sources, with the emphasis on lowering communication complexities between clients and the main central server and training complexities at the client level. \n\nSpecifically, the above is given through the following list of strengths:\n* A head pruning strategy is provided via computation of head importance (that is determined by attention scores).\n* An aggregation mechanism where updates from each head are weighted according to their importance.\n* A client selection strategy based on the loss difference between the global model and the local model at clients.\n* The authors then conducted an extensive experimental setting showcasing the efficacy and efficiency of their proposed approach.\n\nWhile the above is impressive, the paper suffers from the following pair of weaknesses:\n* The current approach is only limited to PEFT approaches, their trainable parameters can be associated with the heads of the multi-head attention mechanism.\n* The paper needs some clarifications associated with some of its parts; see \"Requested Changes\" section.", "claims_and_evidence_yn": "Yes", "claims_explanation": "While the paper has not given any theoretical convergence analysis, however, the paper provides an extensive set of experiments showcasing the efficacy and efficiency of the proposed approach. The intuition behind pruning heads is backed by a previous paper, and the experiments also validate the claim that pruning heads can lead to almost no drop in accuracy.\n\nThe authors also dissected the proposed approach into its contributing elements and assessed their effect on the approach, showcasing the need for their combination. Finally, the authors extended their ablation study to other backbone models to show that their approach is general and can be applied to any MHA model.", "audience_yn": "Yes", "audience_explanation": "The implication of this work is linked with running MHA models on low-resource devices, such as mobile phones, which gives a step towards extending the applicability and availability of such models to all users, knowing or oblivious to such technologies.", "requested_changes": "In what follows, I will list questions and comments concerning the paper:\n* In Equation 6: For the max operation to be valid here, the multiplication here is not the implied dot product, right?\n* In Equation 7: How do you decide the threshold? How is it computed?\n* In Equation 9: Wouldn't it be better to take max over abs?\n* In Table 3: How come that FedAvg has different results? By definition, this approach has no head pruning.", "broader_impact_concerns": "There are no concerns on the ethical implications of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OZpOswschV", "reviewer": "Reviewer_g25m", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes leveraging multi-head attention structure in federated learning to improve training and communication efficiency by pruning and importance-weighted aggregation of attention heads.\n\n---\n### Strengths:\n- The paper presents a clear motivation for addressing federated learning challenges in NLP, and the writing is mostly easy to follow.\n- The paper is clear about the limitations of FL in the field of NLP and recommends an interesting approach to solve the same.\n- The inclusion of detailed ablation studies helps to justify the design choices proposed by the authors.\n\n---\n### Weaknesses:\n- The statement (in Section 2) that prior works have not exploited the multi-head attention mechanism for federated learning is not entirely accurate. Recent studies such as FedTP [1] and FeDeRA [2] have indeed explored this aspect. Acknowledging and contrasting these works would strengthen the paper’s positioning and contribution.\n- The use of the maximum query-key dot product as a proxy for attention head importance may be problematic. Attention scores, especially maximum values, can be noisy and do not always correlate well with actual head contribution to model performance. Several studies [3] have discussed the limitations of using attention weights as importance metrics. It would benefit the paper to provide a thorough theoretical justification for this choice to better support the importance metric & consequently the pruning strategy.\n- The pruning approach currently excludes value vectors in its importance calculation (Equation 6), which may omit crucial information about the head’s output. Can the authors justify why the value vectors were left out? Additionally, the method for determining the pruning threshold is unclear. \n- The baselines provided are somewhat limited. Incorporating comparisons with additional recent works [1, 2, 4] and more advanced aggregation algorithms like FedProx would provide a more comprehensive evaluation of the proposed method’s effectiveness.\n- The comparisons presented in Tables 2, 3, and 4 raise some fairness concerns. Specifically, the proposed method is compared against a baseline using random client selection and no sparsity, while a stronger baseline, FedAvg with loss-based client selection, reportedly outperforms the proposed approach by about 2% as per Table 1. Including this stronger baseline in the comparative tables would help clarify the method’s advantages and limitations. Additionally, can the authors justify the claim that this baseline represents an \"upper bound\"?\n\n\n---\n[1] Li et al. FedTP: Federated Learning by Transformer Personalization. arXiv: 2211.01572\n\n[2] Yan et al. FeDeRA: Efficient Fine-tuning of Language Models in Federated Learning Leveraging Weight Decomposition. arXiv: 2404.18848.\n\n[3] Jain et al. Attention is not Explanation. EMNLP 2019.\n\n[4] Kuo et al. Federated LoRA with Sparse Communication. arXiv:2406.05233", "claims_and_evidence_yn": "No", "claims_explanation": "The experiments and ablation studies generally support the theoretical claims made in Section 3.\n\nHowever, while the experimental setup is reasonable, the chosen baselines do not adequately represent the current state-of-the-art. Although the selected baselines are plausible, including additional well-established methods is necessary to fully assess the novelty and effectiveness of the proposed approach. Furthermore, some experiments and results appear to lack critical comparisons, even within the limited set of baselines provided.", "audience_yn": "Yes", "audience_explanation": "Yes, definitely. Researchers within TMLR’s audience would likely find the paper’s approach to leveraging multi-head attention quite interesting.", "requested_changes": "Critical changes are already discussed in Weaknesses. \n\n- The paper’s citations are not properly formatted, which significantly impacts readability. We kindly request the authors to address this issue to enhance the presentation quality.\n- Considering recent studies have demonstrated that more sophisticated pruning strategies, such as selectively avoiding pruning in certain layers, can lead to improved performance, a more thorough discussion of these strategies would strengthen the paper. At minimum, these aspects should be acknowledged as avenues for future research.\n- Additionally, it would be beneficial for the authors to elaborate on their decision to use synchronous FL and whether asynchronous approaches were explored or considered. Given the practical advantages of asynchronous aggregation in federated and distributed ML, a discussion referencing relevant literature on asynchronous strategies would add valuable context.", "broader_impact_concerns": "No broader impact statement is necessary.", "additional_comments": "N/A", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)