Towards Federated Low-Rank Adaptation of Language Models with Rank Heterogeneity (2025)

Abstract: Low-rank adaptation (LoRA) offers an efficient alternative to full-weight adaptation in federated fine-tuning of language models, significantly reducing computational costs. By adjusting ranks for each client, federated LoRA enables flexible resource allocation. However, we observe that heterogeneous ranks among clients lead to unstable performance. Our analysis attributes this instability to the conventional zero-padding aggregation strategy, which dilutes information from high-rank clients during model aggregation. To address this issue, we propose a replication-based padding strategy that better retains valuable information from clients with high-quality data. Empirically, this approach accelerates convergence and enhances the global model’s predictive performance.

## introduction modern language models have shown unprecedentedly strong performance on many tasks @xcite , but they also have unprecedentedly many parameters. their gigantic sizes become especially problematic in federated fine-tuning of language models, where the cost to compute and communicate local gradients grows proportionally to the number of parameters @xcite . to address this, recent works adopt low-rank adaptation (lora; @xcite ) for federated fine-tuning of language models. instead of tuning all weights, lora freezes original weights and trains only the update parametrized as a product of two low-rank matrices. this reduces the number of parameters, thus reducing the computation and communication needed @xcite . a key promise of federated lora is its potential to improve the resource-accuracy tradeoff by adjusting client-wise ranks @xcite . such rank-heterogeneity provides not only a handy way to tune client-wise computation and communication budgets, but also a mean to bias the global update toward certain clients that are considered giving higher-quality gradient estimates. in this work, we identify a critical shortcoming of existing rank-heterogeneous federated lora methods for language models. whenever the quality of clients varies significantly, conventional rankheterogeneous lora struggles to converge faster than naïve rank-homogeneous lora. our analysis suggests that such underperformance might be due to suboptimal aggregation strategy; to aggregate lora updates with disparate ranks, typical works adopt zero-padding strategy, i.e., matching the dimensionality by appending all-zero rows and columns to the low-rank-decomposed parameter updates @xcite . this strategy may not be optimal whenever there exists some clients which provides much higher-quality information, as the information from such clients can be made less relevant by being averaged with padded zeros. to tackle this problem, we develop a simple yet effective fix, called replication strategy. to avoid having highly relevant information from being diluted, we pad lower rank updates with rows and columns replicated from high-priority clients, instead of zeros. empirically, the proposed method achieves faster convergence to the higher accuracy than existing rank-homogeneous and heterogeneous paradigms. in short, our contributions are: • we identify the shortcomings of existing rankheterogeneous federated lora frameworks for language models, i.e., unexpected slow convergence under high client quality disparity. lora is a parameter-efficient fine-tuning (peft) method that keeps the pretrained weights fixed and only trains newly added parameters @xcite . more concretely, consider fine-tuning a pretrained weight matrix w @xmath0. lora reparametrizes the updated weight matrix w @xmath1: federated lora. in federated lora with k clients, the server receives k different lora updates from the clients. that is, the server receives in rank-homogeneous lora (i.e., r @xmath2, a basic way to aggregate the updates from the clients may aggregated by taking an average for both @xcite . concretely, one performs the aggregated lora weights are then distributed to each client, which is updated further locally until the next communication round. ## shortcomings of the zero-padding our first observation is that the rank-heterogeneous federated lora with zero-padding tends to perform worse than rank-homogeneous lora, whenever the dataset quality varies significantly over the clients (will be shown later in section 6, figure 2 ). here, we have varied the dataset quality of each clients by drawing local data from dirichlet distribution, as in @xcite . here, the client with larger and more balanced datasets are considered of higher quality, as they achieve higher local accuracy during the early training. we have assigned higher ranks to the higher-quality clients. why can zero-paddings hurt? we hypothesize that such unexpected underperformance of zeropadding is due to the fact that padded zeros tend to dilute useful information captured by high-quality clients. to see this, consider averaging k weight matrices a 1 , ã2 , . . . , ãk where a 1 is of rank r 1 and ãi are of rank r 2 < r 1 , which is zero-padded with r 1r 2 all-zero rows. by averaging, the top r 2 rows may retain the same relative scale as the original weight. however, the remaining r 1r 2 rows may have the relative scale of 1 /k, having their impact on the overall model much diminished as the number of clients grow. indeed, our empirical analysis supports this hypothesis; ## method: replication strategy to address this shortcoming, we develop a very simple yet effective method, called replication strategy. instead of padding all-zero vectors, we replicate the rows and columns from the high-rank clients and append them to low-rank clients (figure 1 ). concretely, we first consider a simple case where we have one high-rank client and one low-rank client; let ∆w 1 = b 1 a 1 be the high-rank parameter updates from the first client with some rank r 1 , and let ∆w 2 = b 2 a 2 be the low rank parameter update from the second client with rank r 2 < r 1 . then, the 'row/column-replicated' version of the low rank matrix is given by we emphasize that the overall communication cost remains unchanged. since the replication process is performed exclusively on the server, we can enjoy the advantages of our method without any additional communication overhead. mechanism for allocating high-rank. instead of manually inspecting local datasets to see which client has a high quality dataset (and thus high rank should be allocated), we adopt a simple loss-based criterion to assign high rank. first, we allocate low rank to all clients. after the first local update phase, the server select top-k clients with the highest validation accuracy, and allocate a high rank. ## results the experimental results are given in figure 2 . the leftmost data point denotes the accuracy at initialization (thus can be ignored when comparing baselines), and the subsequent data points denote the test accuracies after each communication round. distilbert. (left two) we first observe that the proposed replication strategy (red) achieves the fastest convergence over all methods in both cases. in particular, the strategy closely achieves the peak test accuracy in two communication rounds. in terms of the final accuracy, the proposed strategy is also among one of the best, together with the communication-heavy option (homogeneous rank 20; orange) which only slightly outperforms on ag's news. zero-padding strategies (dotted lines with circles) converge slower than rankhomogeneous options, with frobenius padding converging slightly faster than naïve. among rankhomogeneous models, the one with a higher rank tends to converge faster to a higher final accuracy. albert. (right two) similarly, our method achieves a the fastest convergence to the high accuracy, only slightly worse than the communicationheavy case (homogeneous rank 20). in ag's news, the homogeneous lora tend to perform slightly better than the replication-based padding after the very first round; this is because the quality of the high rank client selected in the step by our method happened to be worse than other high rank clients. however, our method quickly starts to outperform the baselines in the subsequent rounds; this suggests that our method performs robust w.r.t. the suboptimalities in the high rank client selection. ## conclusion we have identified and analyzed the drawbacks of the zero-padding method during the aggregation process when using heterogeneous lora in federated fine-tuning of language models and proposed a replication-based padding method to address these issues. we have experimentally demonstrated that this method achieves faster convergence with lower resource usage compared to homogeneous lora with high ranks. this suggests that assigning higher ranks to only a limited set of clients-while leaving others with lower ranks-can better align with client resources and data, optimizing overall performance. additionally, this study focuses on a single high rank and a single low rank, allowing for exploration of multiple ranks to better manage resource and data heterogeneity. we believe that our research opens up new challenges and opportunities in federated fine-tuning, and we are confident that this study will contribute to more efficient learning.