FLoRA: Federated Fine-Tuning Large Language Models with Heterogeneous Low-Rank Adaptations

Authors: ziyao  wang
zheyu  shen
yexiao  he
guoheng  sun
hongyi  wang
lingjuan  lyu
ang  li
  genbio
  ai
  sony reasearch
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 19:49:26

Abstract: The rapid development of Large Language Models (LLMs) has been pivotal in advancing AI, with pre-trained LLMs being adaptable to diverse downstream tasks through fine-tuning. Federated learning (FL) further enhances fine-tuning in a privacy-aware manner by utilizing clients' local data through in-situ computation, eliminating the need for data movement. However, fine-tuning LLMs, given their massive scale of parameters, poses challenges for clients with constrained and heterogeneous resources in FL. Previous methods employed low-rank adaptation (LoRA) for efficient federated fine-tuning but utilized traditional FL aggregation strategies on LoRA adapters. This approach led to mathematically inaccurate aggregation noise, reducing fine-tuning effectiveness and failing to address heterogeneous LoRAs. In this work, we first highlight the mathematical incorrectness of LoRA aggregation in existing federated fine-tuning methods. We introduce a new approach called FLoRA that enables federated fine-tuning on heterogeneous LoRA adapters across clients through a novel stacking-based aggregation method. Our approach is noise-free and seamlessly supports heterogeneous LoRAs. Extensive experiments demonstrate FLoRA's superior performance in both homogeneous and heterogeneous settings, surpassing state-of-the-art methods. We envision this work as a milestone for efficient, privacy-preserving, and accurate federated fine-tuning of LLMs.

Main paper: 
 introduction: 
 The Large Language Models (LLMs) have shown remarkable performance on various tasks, such as chatbots [1] , virtual assistants [4] , search engines [11] , and healthcare [20; 18] . However, adapting pretrained LLMs (e.g., Llama 2 [22] ) to downstream tasks requires tremendous computation resources to fine-tune all the model parameters. To mitigate this issue, a variety of parameter-efficient fine-tuning (PEFT) methods have been proposed. One of the most widely used PEFT methods is low-rank adaptation (LoRA) [10] . As shown in the top of Figure 1 , LoRA adds a parallel branch of trainable adapters A and B to compute the model update ∆W, where the ranks of A and B are much smaller than the pre-trained model parameter W. When applying LoRA for fine-tuning, only A and B are updated while the entire W is frozen, thereby significantly reducing the GPU memory consumption.Fine-tuning LLMs requires ample data for adaptation to specific downstream tasks [13; 7] . Often, this data is dispersed across a multitude of devices, raising privacy concerns. For instance, aggregating medical data from hospitals for centralized LLM fine-tuning poses significant challenges. Consequently, to facilitate fine-tuning without compromising private data, federated learning (FL) becomes essential, enabling LLM fine-tuning across distributed clients while preserving data privacy [16; 27; 33; 24] . In this work, we focus on federated fine-tuning, enabling distributed clients to collaboratively fine-tune LLMs for adaptation to downstream tasks while preserving data privacy.Prior work, FedIT, proposed a federated fine-tuning method [28] , integrating LoRA with FedAvg [16] . In each FL round of FedIT, clients fine-tune LoRA modules using their local data and then send the fine-tuned modules to the server. The server averages all the local LoRA modules to obtain a global LoRA. Since only the weights of the LoRA modules are fine-tuned and communicated, FedIT effectively reduces both computation and communication costs.However, FedIT faces two key issues. First, the naive averaging of local LoRA modules in FedIT introduces noise to the global model update. Specifically, FedIT averages local A and B independently, which introduces mathematical errors to the global LoRA. In short,The cause of aggregation noise:A × B FedIT ̸ = A × B mathematically correct.We will elaborate on this issue in Section 2 with theoretical analysis. Such an inaccurate aggregation will hinder convergence, leading to higher fine-tuning costs. Second, due to the heterogeneous data distribution [31; 12] and heterogeneous hardware resources, clients need to adapt LoRA ranks [30] according to the system and data heterogeneity. However, FedIT cannot aggregate local LoRAs with heterogeneous ranks.Homogeneous OnlyFedIT: In this work, we present FLORA, an aggregation-noise-free federated finetuning method that supports heterogeneous LoRAs. Specifically, as shown in Figure 2 , we propose to stack the local LoRA modules A k and B k separately to construct the global LoRA modules A and B, where A k and B k denote the corresponding LoRA modules on the k-th client. This stacking method is theoretically proven to be accurate for the aggregation of local LoRA modules (Section 3.1). Additionally, it can naturally accommodate heterogeneous LoRA settings (Section 3.2), since stacking does not require the local LoRA modules to have identical ranks across clients. The noise-free aggregation of FLORA accelerates convergence, which will in turn improve the overall computation and communication efficiency of federated fine-tuning. Furthermore, FLORA can effectively cater to heterogeneous data and computational resources across clients, where heterogeneous ranks are applied. Our key contributions are summarized as follows:Averaging• We propose FLORA, a federated fine-tuning algorithm based on LoRA that can perform noise-free aggregation of local LoRA modules. Theoretical analysis shows that FLORA eliminates the meaningless intermediate term in the global model update, leading to faster convergence and improved performance. • The proposed stacking mechanism for aggregating LoRA modules supports heterogeneous LoRA ranks across clients, accommodating data and system heterogeneity in realistic settings. This encourages the broader participation of clients with heterogeneous data and resources in federated fine-tuning. • We use FLORA to fine-tune LLaMA, Llama2 [21] and TinyLlama [29] on four benchmarks for two downstream tasks. Results show that FLORA surpasses state-of-the-art methods for both homogeneous and heterogeneous settings.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The paper has a clear abstract and an introduction section.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
preliminaries: 
 Fine-tuning LLMs with LoRA. LoRA [10] uses two decomposed low-rank matrices to represent the update of the target module: where W ∈ R m×n and W ′ ∈ R m×n denote the pre-trained and fine-tuned parameters of target modules (e.g., attention modules), respectively. A and B are low-rank decomposition of ∆W, where A ∈ R r×n , B ∈ R m×r , such that ∆W = BA with the identical dimensions as W and W ′ . The rank of LoRA, denoted by r, is typically significantly smaller than m and n, leading to dramatic parameter reduction of ∆W. During the fine-tuning phase, LoRA optimizes matrices A and B instead of directly updating W, thus achieving substantial savings in GPU memory usage. For example, in the context of the LLaMA-7B [21] , the original dimension of attention modules is W ∈ R 4096×4096 , setting the LoRA rank to 16 reduces the decomposed matrices to A ∈ R 16×4096 and B ∈ R 4096×16 . This approach decreases the number of trainable parameters to merely 0.78% of the entire parameter of the pre-trained model, offering a significant GPU memory footprint reduction.EQUATIONFedIT: Averaging Homogeneous LoRA. The most widely used FL algorithm, i.e., FedAvg [16] , aggregates all the local model updates by weighted averaging to update the global model in each communication round:EQUATIONwhere W ′ and W denote the global model parameters before and after a communication round. ∆W k represents the local model update from the k-th client, with p k being the corresponding scaling factor that is typically weighted by the local data size, and ∆W represents the global model update.FedIT [28] directly integrates FedAvg with LoRA to enable federated fine-tuning, where each client fine-tunes LoRA modules with a homogeneous rank. Specifically, the clients download the pre-trained LLM from the server, locally initialize and fine-tune the LoRA modules, and then send the updated LoRA modules to the server. The server updates the global LoRA modules A and B by independently applying weighted averaging across all local modules A k and B k :EQUATIONThis aggregation of FedIT is almost the same as FedAvg except that only the LoRA modules are trained and communicated. However, such a naive aggregation introduces additional issues for federated fine-tuning. First, each single module A or B is not the model update, and only BA represents the model update. Thus, averaging A k and B k independently to compute the aggregated gradients will introduce noises to the global model update. Here we use a simple example to explain how the noise is generated, and we assume that two clients are applying FedIT to perform federated fine-tuning. In a communication round, the two clients train A 0 , B 0 and A 1 , B 1 respectively. The local model updates ∆W 0 and ∆W 1 are the product of corresponding LoRA modules:EQUATIONAccording to Equation 2, the expected global model update ∆W can be obtained by weighted averaging ∆W 0 and ∆W 1 :EQUATIONHowever, according to Equation 3, FedIT aggregates A and B independently:EQUATIONThe global model update in Equation 6differs from the expected one in Equation 2 illustrates, the averaging algorithm in FedIT is an inaccurate aggregation method, leading to slower convergence and higher computation costs.The other deficiency of FedIT is that it cannot support aggregation on heterogeneous LoRA modules. The local data in FL may exhibit significant heterogeneity across clients [31; 12] . If a client configures a higher rank than the actual one required by the local data complexity, this may result in overfitting. Conversely, if the rank is too small, it may lack the necessary generalization capacity to effectively learn from the local dataset (Figure 4 ). Moreover, the heterogeneous computational resource across clients also requires heterogeneous rank deployment, e.g., clients with smaller memory can only afford to train LoRA modules with smaller ranks. AdaLoRA [30] has been proposed to adapt LoRA ranks based on available computation resources. Therefore, deploying heterogeneous ranks across clients is a pressing requirement for accommodation to data and system heterogeneity. However, according to Equation 3, FedIT is only able to aggregate LoRA modules with the homogeneous rank.3 Proposed Method: FLoRA 1), the element at position (x, y) of the model update ∆W is computed as the sum of the products of corresponding elements from the x-th column of B and the y-th row of A:EQUATIONwhere δ xy represents the element at position (x, y) in ∆W. a yi , b xi are the elements at positions (y, i) and (x, i) in A and B, respectively. According to Equation 3.1, the model update in LoRA can be expressed as the sum of the products of the corresponding rows of A and the columns of B.To illustrate this concept further, let us consider a simplified example where the dimensions of LoRA modules are given by A ∈ R 2×3 and B ∈ R 3×2 . As described in Equation 8, A and B can be decomposed to two sub-matrices with rank r = 1, and the product of A and B then are computed as the sum of the products of two respective sub-matrices: To address the aggregation challenge from an alternative perspective, let us consider the scenario where we have multiple pairs of LoRA modules, A k , B k , optimized by the clients. Each pair satisfies the dimensions A k ∈ R r k ×n and B k ∈ R m×r k . Similar to Equation 8, the sum of the products of these module pairs is the product of the stacked modules, i.e., K k=1 B k A k = BA, where B represents the stacking of all B k modules aligned through dimension m and A is the stacking of all A k aligned through dimension n. Figure 2 visually illustrates this concept, where the orange, green, and blue rectangles symbolize A k , B k , and their respective products. The aggregation of three products mirrors the product of the stacked B and A from all B k and A k pairs trained by clients. This mechanism demonstrates that, in the context of federated fine-tuning, we can achieve a noise-free aggregation of local updates by simply stacking the local LoRA modules. This process also avoids transmitting the full model parameters, thus reducing communication costs.To facilitate our discussion, we introduce the stacking operation symbolized by "⊕" to denote the module aggregation as depicted in Figure 2 . This operation is mathematically defined as: In Equation 9, "⊕" indicates that for A, each subsequent module is vertically stacked below the preceding one, whereas for B, each module is horizontally stacked to the right of the one before it.EQUATIONWe can now formalize our conclusion regarding the aggregation of LoRA modules. The sum of the products of K LoRA module pairs is equivalent to the product of their stacked matrices:EQUATIONThis foundational principle will guide the design of FLORA, as it allows for the efficient and effective aggregation of local updates without the transmission of entire model parameters.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The Limitation section is presented in the Appendix.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
flora: stacking-based federated fine-tuning for heterogeneous lora: 
 The stacking-based aggregation facilitates not only the accurate aggregation of LoRA modules but also inherently supports the heterogeneous LoRA ranks. This approach imposes no constraints on the ranks of each local LoRA module as long as each client fine-tunes the same pre-trained model, i.e., they share the same dimension m and n.By employing the stacking-based aggregation mechanism, we introduce FLORA, an approach designed to facilitate federated fine-tuning of LLMs with heterogeneous LoRA. Let us use a concrete example to illustrate the key steps of applying FLORA, where K heterogeneous clients are involved in fine-tuning an LLM, and the pre-trained parameters are denoted by W.Initialization. The server first disseminates the pre-trained model parameters W to all K clients. Then, the clients initialize their local LoRA modules based on the complexity of local data and available local resources. The adaptation of LoRA ranks is beyond the scope of this paper, but existing work like AdaLoRA [30] can facilitate the rank adjustment.Local Fine-tuning. Following initialization, the clients train their local LoRA modules with the local data for several iterations. Then, the clients send the local LoRA modules back to the server. Note that the clients initialize local LoRA modules each round before local fine-tuning.Stacking-based LoRA Aggregation. Upon receiving the heterogeneous LoRA modules from participating clients, the server proceeds to aggregate them by stacking all B k and A k according to Equation 10, resulting in the global A ∈ R ( K k=0 r k )×n and B ∈ R m×( K k=0 r k ) . The aggregation process of FLORA can be described as follows:EQUATIONwhere p k represents the scaling factor for each local update, determined by the relative size of the local data to the global data:EQUATION)Note that the scaling factor p k should be only applied to one of A k and B k to avoid squaring the factor in the final model update BA. This method ensures a noise-free aggregation mechanism as described in Equation 10.Update Local Models. After each round of noise-free aggregation, the server redistributes the updated global LoRA modules A and B back to the clients. The clients then proceed to update the local models using BA and continue the fine-tuning. Using the stacking approach, the dimensions of updated global LoRA modules A and B are larger than those of FedIT, potentially leading to larger communication overhead in each round. However, empirical observations indicate that federated fine-tuning typically requires only a limited number of communication rounds to achieve satisfactory results, as detailed in Section 4. In addition, it is important to note that the LoRA modules A and B constitute a small fraction of the overall size of the pre-trained model, which is distributed to clients during the initialization phase. Thus, the additional communication overhead of the stacking approach is negligible and does not significantly impact the efficiency of federated fine-tuning.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Answer: [Yes] Justification: The code is on GitHub. We already have a discussion with other researchers about the code.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experiments: 
 The key features of FLORA are (i) noise-free aggregation and (ii) support for heterogeneous LoRA modules. In this section, we verify these key features across various LLM fine-tuning tasks. We first study the performance of FLoRA and compare it against FedIT under homogeneous settings to demonstrate the advantages of noise-free aggregation [28] . Then, we examine performance in a synthetic heterogeneous setup and compare FLORAwith a vanilla zero-padding method. Finally, we conduct ablation studies on the scaling factor, the heterogeneity of LoRA ranks, and the extra communication overhead of FLoRA.Models, Datasets and Experiment Settings. We employ three Llama-based models with different scales in our experiments: TinyLlama with 1.1 billion parameters [29] , and the 7 billion parameter versions of Llama [21] and Llama2 [22] , evaluating FLORA across different model capacities. Following the configurations in the original LoRA paper [10] , the LoRA modules are applied to the self-attention layers only.We use the Databricks-dolly-15k [28] instruction dataset, Alpaca dataset [19] , and Wizard dataset [14] for the question-answering (QA) task, and Wizard and ShareGPT for the chat assistant task. We evaluate the federated fine-tuned models on MMLU [8] for the QA task and MT-bench [32] for the chat assistant task, respectively. We sample 10 clients uniformly at random following the non-IID setting in FedIT [28] . The other experimental configurations are elaborated in Appendix A.Baselines. We compare FLORA with four baselines. (1) FedIT: It is the SOTA federated finetuning method [28] that integrates LoRA with FedAvg. We only apply FedIT to homogeneous LoRA experiments as it does not support heterogeneous LoRA. (2) Zero-Padding: It is an approach that enables FedIT to support heterogeneous LoRA [3] . It extends all the heterogeneous local ranks to the maximum rank among the clients and pads their remaining parts by 0. (3) Centralized Fine-tuning: we compare FLORA with centralized LoRA with the same hyperparameters and configurations. (4) Standalone: the client fine-tunes the pre-trained model locally without federations. Homogeneous LoRA. We first evaluate the performance of FLORA with homogeneous LoRA. Specifically, all the clients share the identical LoRA rank of 16. As Table 1 depicts, FLORA achieves consistently better performance than FedIT across all the evaluated models and tasks. This is evident in the MT-bench scores for both TinyLlama and Llama models, where FLORA's performance exceeds that of FedIT by at least 0. As Table 1 and Table 4 illustrate, FLORA not only adapts to heterogeneous ranks without performance degradation but also maintains consistency with the results observed in most homogeneous settings. This contrasts sharply with the performance of FedIT, where the application of zero-padding significantly degrades its performance on MMLU and MT-bench. It reveals that zero-padding exacerbates FedIT's inherent noise issues in the aggregation process, posing significant challenges in managing fine-tuning performance. For example, by applying the zero-padding method, the MMLU accuracy of Llama model fine-tuned with Alpaca dataset dramatically drops to 7.97%. The results demonstrate that FLORA not only accommodates heterogeneous LoRA ranks effectively but also sustains robust training performance compared to baseline methods. It efficiently facilitates the participation of devices with varied computational capacities in heterogeneous federated fine-tuning tasks. Additionally, FLORA can be seamlessly integrated with AdaLoRA [30] , which dynamically adjusts the LoRA rank on the clients, the results are presented in Appendix A. The Impact of Scaling Factor. The scaling factor, denoted as p k in Equation 12, plays a pivotal role in the efficacy of FL [23] . We conduct experiments investigating how varying scaling factors influence the performance of FLORA.Given that the default scaling factor is set to 0.1 for all clients, assuming 10 clients with equal local dataset sizes as per Equation 12, we explored the effects of alternative scaling factors, namely 0.01, 0.05, and 0.2. The results are summarized in Figure 5 . The results do not reveal a clear pattern or optimal scaling factor for federated fine-tuning across different settings. The efficacy of a specific scaling factor appears to be contingent upon the dataset, task, and model in use. For example, when fine-tuning TinyLlama on the Dolly dataset, a lower scaling factor of 0.01 yields the highest accuracy, significantly outperforming the 0.1 and 0.2 scaling factors. Conversely, the model fine-tuned on Wizard dataset demonstrates a preference for a higher scaling factor of 0.2, achieving the best performance, whereas the lowest scaling factor of 0.01 was the least effective. In the case of the Llama model, larger scaling factors consistently facilitated better fine-tuning performance. Applying FLORA to Dolly and Alpaca shows the optimal performance with a scaling factor of 0.2. These observations suggest that the choice of an appropriate scaling factor is highly dependent on specific datasets and model characteristics, underscoring the necessity for a tailored approach in federated fine-tuning.The Impact of Heterogeneous LoRA Ranks. Although the above results demonstrate FLORA effectively enables the federated fine-tuning with heterogeneous LoRA, it is worth further investigating how the federated fine-tuning improves the local models with various ranks. Motivated by this, we evaluate MT-bench scores for local models with LoRA ranks of 64, 32, 16, 8, and 4, presenting the results in Regarding the LoRA rank's impact, a rank of 8 consistently yields strong performance across various models and datasets. However, performance diverges at extreme ranks; for instance, the TinyLlama model fine-tuned on Wizzard with the LoRA rank of 64 underperforms the ones with smaller ranks, but the Llama model with the rank of 64 excels the counterparts with smaller ranks. This also demonstrates the heterogeneous rank deployment across clients is a realistic setting. These observations suggest a potential positive correlation between optimal LoRA rank and model capacity, motivating further exploration in future research. 
discussion: 
 The Communication Overhead of FLORA. As discussed in Section 3, the server needs to send global LoRA modules to the clients in FLORA, potentially raising concerns about increased commu-nication overhead. To quantify this, we compare the communicated parameters of full fine-tuning, FedIT, and FLORA over three communication rounds.Figure 6 : The ratio of communicated parameter numbers to full fine-tuning.As Figure 6 shows, although FLORA transmits slightly more parameters than FedIT, it still significantly reduces the overhead compared to full fine-tuning. This is due to the fact that the primary communication load in federated fine-tuning, especially with large models, is the initial full model parameter transmission. Subsequent rounds primarily involve smaller updates (e.g., LoRA matrices). Thus, even though FLoRA introduces additional communication for these updates, the overall impact on total communication costs remains marginal, making it comparable to FedIT's costs. Despite the minor communication increase compared to FedIT, FLORA enhances fine-tuning effectiveness and supports heterogeneous LoRA ranks, making it a preferable solution in federated fine-tuning.The Privacy Preservation of FLORA. The requirement of FLORA to stack the LoRA modules uploaded by all clients introduces a potential privacy concern, as malicious clients might infer the LoRA matrices of other clients through the global LoRA modules sent from the server. To address this issue, we split all the local LoRA modules into sub-modules with rank=1 and then stack the sub-modules together in random order. This approach prevents malicious clients from recovering the local LoRA modules from other clients. In addition, FLORA is also compatible with standard privacy mechanisms such as encryption [2] and differential privacy (DP) [25] , aligning it with the privacy-preserving nature of FL. 
related work: 
 Parameter-efficient Fine-tuning of LLMs. Parameter-efficient fine-tuning (PEFT) aims to reduce the number of trainable parameters. BitFit [26] fine-tunes only the biases while achieving similar accuracy with full fine-tuning. Other works such as [9] and [17] apply transfer learning that adds pre-trained adapter layers between transformer blocks. LoRA [10] adopts the product of two low-rank matrices to represent the gradient in full fine-tuning, which achieves memory-efficient fine-tuning. AdaLoRA [30] optimizes LoRA by adaptively allocating the parameter budget, which enhances the flexibility of LoRA. There are also many works regarding optimizing LoRA in various aspects [5; 6; 15] . Federated Fine-tuning of LLMs. Federated fine-tuning aims to extract knowledge from multiple ondevice datasets while preserving data privacy. FedIT [28] leverages the FL framework for fine-tuning LLMs. It uses LoRA as the local fine-tuning strategy. However, concerns related to the deficiency in supporting heterogeneous LoRA limit its utilization. [3] tries to solve this problem by zero-padding the local LoRA modules. However, this padding process causes additional computing overhead. Besides, it separately averages A and B modules, introducing noise to the global model.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Details are in the Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 In this work, we identified the limitations in current federated fine-tuning methods (e.g., FedIT), and the challenges of applying federated fine-tuning in realistic settings, i.e., the heterogeneous LoRA ranks across clients. To overcome these practical challenges and broaden the applicability of federated fine-tuning, we introduced FLORA to enable the accurate aggregation on heterogeneous LoRA modules using the proposed stack-based LoRA aggregation mechanism. Our extensive experiments demonstrate that FLORA outperforms the SOTA method in both homogeneous and heterogeneous LoRA settings. Moreover, our inspiring results provide valuable insights for future research in the federated fine-tuning of large language models in a lightweight and accurate manner. Alpaca dataset. The Alpaca dataset contains 52K instruction-following data used for fine-tuning the Alpaca model [19] . This dataset is believed to be diverse enough for fine-tuning LLMs.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: The results are reproducible with low errors.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: It is only a technical paper, which has no social impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: All the models and datasets used are open-sourced.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: Models and datasets are cited.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Wizard dataset. The Wizard dataset we use is the training data of the WizardLM model. It includes 70k pairs of instructions and outputs. The Wizard dataset generally features more complex instructions compared to the other datasets. Its fine-tuning results are typically better, which has been confirmed by our experiments, especially those evaluated by the MT-bench scores.ShareGPT dataset. The ShareGPT dataset is a collection of approximately 52,000 conversations scraped via the ShareGPT API. The conversations in ShareGPT include both user prompts and responses from ChatGPT. In our experiments, we split the conversation dataset into question-answering pairs. MMLU test set. The MMLU dataset is a widely used question-and-answer dataset in LLM finetuning. It has 14,024 questions in 57 different subjects, which can evaluate the logical reasoning capabilities of LLMs. We selected 1444 samples from the dataset for a quick and comprehensive evaluation.MT-bench evaluation. MT-bench is a set of challenging multi-turn open-ended questions for evaluating chat assistants [32] . It evaluates the performance of LLMs by using the GPT-4 API to score the LLM-generated conversations. LLMs that behave more like GPT-4 will receive higher scores.In all our experiments, the learning rate of fine-tuning is set to 0.0003; the batch size is 128 and the micro batch size is 16. Due to the large dataset and model sizes selected, federated fine-tuning consumes significant computational resources and time. Therefore, we opted for fewer fine-tuning rounds (even just one round) to ensure that we could observe enough data. Additionally, the MMLU dataset is prone to overfitting on these large datasets, resulting in a decrease in accuracy. Therefore, fewer training rounds ensure the effectiveness of the observed phenomena. Table 2 shows the fine-tuning rounds and local epochs we selected. Integrating FLORA with AdaLoRA All the observations about the impact of rank on the model performance, despite being influenced by data heterogeneity, still manage to reveal the importance of selecting an appropriate LoRA rank for a specific task. Thus, some algorithms such as AdaLoRA [30] are designed to adaptively adjust the LoRA rank to optimize the model performance and save computational resources. With our support for heterogeneous LoRA, we can flexibly utilize AdaLoRA with adaptive LoRA ranks. We conducted corresponding experiments to demonstrate that we can use AdaLoRA to further improve the efficiency of federated fine-tuning. We implement AdaLoRA on each client to adjust LoRA modules during local fine-tuning. The results are shown in Table 3 .The "Sum of local ranks" column means the sum of all local LoRA rank values after fine-tuning. Since our FLORA does not adjust the rank, its value is 160, the same as the initial value. On the other hand, AdaLoRA dynamically adjusts the rank to maximize training effectiveness and minimize rank values to save resources. From Table Table 3 , we can see that AdaLoRA on TinyLlama and Llama reduced the sum of local ranks to 120 and 131 from 160 respectively. We further conclude that FLORA+AdaLoRA can further reduce the trainable parameter count while ensuring comparable or even improved performance compared to simply using LoRA on the clients. Our support for such rank adaptation further demonstrates the effectiveness and applicability of the FLORA approach.Experiment results of Llama2. Due to the inherently strong performance of Llama2, the improvement in the QA dataset is not significant. Therefore, we fine-tuned Llama2 using the Wizard and ShareGPT datasets. Overall, Llama2 exhibits similar experimental results to Tinyllama and Llama. Table 4 shows the comparison between FLORA and our baselines. In the homogeneous and heterogeneous settings, the MT-bench scores of Wizard and ShareGPT all surpass those in FedIT and Zero-Padding. As for the impact of scaling factors in Figure 7 , Llama2 has a similar trend to the Llama-7b model, in which higher scaling factors exhibit better fine-tuning performance.In this section, we demonstrate the convergence of FLORA following the standard convergence analysis in [12] . The FedAvg algorithm exhibits convergence to the global optimum at a rate of O(1/T ) for non-IID (independent and identically distributed) data under full client participation. This convergence is based on four assumptions mentioned in [12] : E∥∇F k (W (t) k , ξ (t) k ) - ∆W (t) k ∥ 2 ≤ σ 2 k for k = 1, ..., K, where ξ (t)k is the subset of training data randomly sampled from k-th client.Assumption 4. The expected squared norm of stochastic gradients is uniformly bounded:E∥∇F k (W (t) k , ξ (t) k )∥ 2 ≤ G 2 for all k = 1, ..., K and t = 1, ..., T , where ξ (t)k is the subset of training data randomly sampled from k-th client.For the convergence analysis of FLORA, we introduce an additional assumption 5 tailored to the specific dynamics of LoRA fine-tuning and its relation to traditional SGD-based full fine-tuning: Assumption 5. (Unbiased LoRA Gradient). The updates applied to LoRA modules by each client serve as unbiased estimators of the gradient that would be directly computed on the base model through SGD:B (t+1) k A (t+1) k -B (t) k A (t) k = η (t) ∇F k (W (t) k |ξ (t)k ). Note that we define the model parameter in t-th round by W (t) . Theorem 1. Based on Assumptions 1-5, we choose k = L µ , γ = max{8k, E}. The local learning rate α k r k = 2 µ(γ+t) . Then, we can deduce that the expectation of the fine-tuning error in FLORA can be bounded by:EQUATIONwhere δ (T ) is the fine-tuning error in T -th round. δ (T ) and M are defined as follows:EQUATION)where L, µ, σ k , and G are defined by the assumptions 1-4. Γ is defined byΓ = F * - K k=1 p k F * kfor quantifying the degree of non-iid. This theorem posits that as the number of rounds T approaches infinity, the expectation of the fine-tuning error δ (T ) converges to zero. In contrast, FedIT deviates from the FedAvg model updating rule as depicted in Equation 2, introducing non-gradient noises through its averaging process. Therefore, it fails to achieve convergence at the rate of O(1/T ) While this deviation does not invalidate FedIT's utility in federated fine-tuning, it significantly impairs its convergence rate and overall effectiveness.Our approach has the limitation that the server sends the stacked LoRA modules to the client, thereby increasing the communication costs. We discussed this limitation both theoretically and experimentally in Section 3 and Section 4, respectively. We believe that the increase in communication overhead is acceptable under the premise of improving fine-tuning effectiveness and accelerating convergence. In addition, due to constraints on computational resources and time, we only utilized Llama models in the experimental section. We aim to observe experimental phenomena of different types of LLM federated fine-tuning in future research and derive more general principles.NeurIPS Paper Checklist 


Decision: This paper proposes a federated learning mechanism to train LLMs with LoRA, The proposal is simple, but practical, with its effectiveness demonstrated in experiments. Therefore, we recommend accepting it.

Reviews: review: paper_summary: paper_summary: This paper focuses on developing a novel aggregation strategy for training LLMs in FL with LoRA. Specifically, they identify that separately aggregating two matrices of LoRA module is not identical to aggregating model updates. Therefore, they propose a stack based aggregation strategy that merge matrices before aggregating them.
strength_weakness: strength_weakness: 1.	The novelty is limited. Although the proposed method is effective, the contribution of the stack-based aggregation strategy may be too limited.

2.	The motivation of the stack based aggregation is not clear. Since the computational resource of the server is usually sufficient, it is totally acceptable that recovering the local update and then aggregating them.
questions: questions: no
limitations: limitations: yes
paper_summary: paper_summary: Previous methods using Low-Rank Adaptation (LoRA) for efficient federated fine-tuning may have led to mathematically inaccurate aggregation noise, reducing the effectiveness of fine-tuning and being unable to handle heterogeneous LoRA. The authors analyzed the mathematical inaccuracies in LoRA aggregation in existing federated fine-tuning methods and proposed a new stack-based aggregation method that supports federated fine-tuning across client machines with heterogeneous LoRA adapters. Extensive experiments demonstrate the superior performance of FLORA in homogeneous and heterogeneous environments.
strength_weakness: strength_weakness: 1. Writing needs to be calibrated. The sentences in line 64 and line 70 are repeated.
2.The discussion on accelerating convergence is not indexed in the main text (Appendix B) and lacks experimental validation.
3. More discussion may be needed on the advantages of FLoRA compared to the implementation of 'the aggregation of local update' in Figure 2.
4. The foundation models chosen for the experiment are all from the Llama series. Are there other types of foundation models that can be used for validation?
questions: questions: see the weaknesses above
limitations: limitations: see the weaknesses above
paper_summary: paper_summary: This paper proposes a novel aggregation algorithm FLoRA for low-rank federated fine-tuning. Compared to prior works, FLoRA stacks the LoRA matrices together rather than averaging them, thus avoid the mathematical errors during the aggregation process on the server. The paper provides experiments compared with baselines and discusses about the efficiency and privacy preservation of FLoRA. The paper provides convergence analysis of FLoRA.
strength_weakness: strength_weakness: (1) The paper misses some related works. For example, [1] also talks about LoRA in federated fine-tuning and optimize its efficiency and accuracy. [1] and other possible related works should be considered and added to the paper.
(2) The experiment results presented in Figure 4 show that some single client perform better than the server, which needs further explanation in the paper. Additionally, the subtitle 'The Impact of Heterogeneous...' here does not match the title 'Standalone...' of Figure 4, which needs further polishing.
(3) The paper should explain why it uses Llama rather than SOTA models (e.g., Llama-3).

[1] Sun, Y., Li, Z., Li, Y., & Ding, B. (2024). Improving LoRA in privacy-preserving federated learning. arXiv preprint arXiv:2403.12313.
questions: questions: Please see the weaknesses. The author might need to explain why they are using Llama rather than stronger models like Llama-2 or Llama-3.
limitations: limitations: The limiations are discussed.

Comments: title: Thanks for your response. We further address your concern here.
Looking forward to your further feedback
Thanks for your reply, we clarify the computation and communication cost here
Looking forward to your response
Follow-up on rebuttal
Thanks for your response. We have addressed your concerns here.
comment: Thanks for the response. I agree with the advantage of saving recovery time of the stack-based aggregation. I decide to increase my score to 5. I hope the authors can add this analysis to the revised paper. Besides, the corresponding verification experiments about the recovery efficiency should also be conducted.
>The simple set-based aggregation can also have a total memory usage of ( 2P_{\text{full}} ). In particular, we can compute B_iA_i one by one and add outcomes one by one.

We do agree that computing B_iA_i one by one has a total memory usage of ( 2P_{\text{full}} ). However, as we mentioned earlier, such sequential computation would significantly increase the time required for the recovery process. If the time for the GPU to perform one $BA$ is t, then calculating for $M$ clients one by one will take approximately $Mt$. This is because the GPUs have strong parallel computing capabilities. The time for the GPU to compute a single LoRA matrix multiplication with $r=1$ is similar to the time required to compute it with $r=M$, whereas performing $M$ separate $r=1$ LoRA matrix multiplications takes $M$ times longer than performing a single $r=M$ multiplication.

We also want to recall that the main drawback of using such a set-based aggregation is the difficulty of utilizing LoRA on the base model. The LoRA modules cannot be utilized without merging them into the base model parameters, which limits the implementation of the fine-tuned model.

>I am confused about why the recovering computation of set-based aggregation is (M+1)/2. Considering an extreme case $r=1$, then the computational cost of set-based aggregation is $Kdd$. The computational cost of stack-based computational cost is also $Kdd$

Here, we need to clarify that the computational overhead we are evaluating primarily concerns GPU memory rather than computational complexity. This is because, in large model fine-tuning/inference, the main bottleneck is memory usage. Current works also mainly focus on reducing memory consumption rather than computational complexity, often even increasing complexity to reduce memory usage (e.g., LoRA). Therefore, our evaluation of computational overhead also focuses on memory consumption. As mentioned in Question 1, if the set-based approach calculates all LoRA modules simultaneously, it will increase memory usage to $(1+M)/2$. However, calculating LoRA modules sequentially would significantly reduce the efficiency of code execution.

Thank you for continuing to engage with us, and we look forward to receiving your further feedback.
Thanks for the response. 
I agree that using stack-based aggregation can improve GPU utilization. Yet, I am confused about the other advantages.

1. The simple set-based aggregation can also have a total memory usage of ( 2P_{\text{full}} ). In particular, we can compute B_iA_i one by one and add outcomes one by one.

2. I am confused about why the recovering computation of set-based aggregation is (M+1)/2. Considering an extreme case $r=1, A \in R^{d * 1}, B \in R^{1 * d}$, then the computational cost of set-based aggregation is $K*d*d$         

The computational cost of stack-based computational cost is also $K*d*d$.
Dear Reviewer 4FRM,

The rebuttal deadline is approaching soon, and we have provided detailed responses to your latest concerns. We kindly request that you review our replies at your earliest convenience. If there are any further questions or issues, please let us know. If our responses have adequately addressed your concerns, we would appreciate a higher score.

Thank you for your time and consideration.
Compared to storing the LoRA modules in a set, stack-based methods can significantly reduce computational costs due to the following advantages:

1. **Reduced Memory Overhead**: During the weight recovery process of LoRA, each LoRA module pair is expanded to the full parameter size. In LoRA fine-tuning with \( M \) clients, if the modules are stored in a set, each module must be recovered separately, leading to a total memory usage of \( (M + 1)P_{\text{full}} \), where \( P_{\text{full}} \) represents the full parameter size. This includes the original parameters and \( M \) newly recovered parameters. In contrast, the stack-based method employed in FLoRA allows for the recovery of the full parameters with **only a single matrix multiplication**, resulting in a total memory requirement of just \( 2P_{\text{full}} \). Given that \( M > 1 \) in federated fine-tuning, this approach leads to substantial computational savings.

2. **Optimized GPU Utilization**: Since parallel computing on a GPU is significantly faster than sequential processing, it is more efficient to compute the stacked LoRA modules collectively rather than individually. This approach better aligns with the GPU’s architecture, leading to faster computations and more efficient resource utilization.
To illustrate the benefits of our method in terms of computation and communication costs, particularly in comparison to the baseline FedIT and the set-based method proposed by the reviewer, we provide a table below. For reference, the communication and computation costs of FLoRA are normalized to 1:

=========================================================================================

Method $\quad$ Communication  $\quad$ Recovering Computation  $\quad$ Accurate aggregation $\quad$ Applicable to inference

-------------------------------------------------------

FedIT   $\quad$ $\quad$ P_{full}/P_{LoRA}  $\qquad$  $\qquad$ 1  $\qquad$  $\qquad$ $\quad$  $\qquad$  $\qquad$     No $\quad$  $\qquad$  $\qquad$  $\qquad$ $\qquad$ Yes

Set-saving     $\qquad$ $\quad$   1   $\qquad$  $\qquad$ $\qquad$  $(M+1)/2$  $\ $ $\ $ $\ $  $\qquad$ $\quad$ $\qquad$ Yes $\quad$  $\qquad$  $\qquad$  $\qquad$ $\qquad$  No   

FLoRA (ours)     $\quad$  $\quad$ 1  $\qquad$ $\qquad$ $\qquad$ $\quad$  1   $\quad$ $\quad$ $\qquad$ $\qquad$ $\quad$ $\qquad$  Yes $\quad$  $\qquad$  $\qquad$  $\qquad$ $\qquad$Yes

=========================================================================================

According to this table, our approach generally outperforms the baselines in important metrics.
Thanks for the reply. Can the stack-based aggregation save the computational cost?
It would be better if the analysis regarding the computational and communication costs could be provided and compared with the aggregation simply using a set.
Dear Reviewer,

We wanted to inform you that we have addressed the additional concerns you raised. As the rebuttal phase is coming to an end, we kindly request your prompt feedback on our rebuttal.

We appreciate your time and consideration and look forward to your response.

Best regards,
Authors
Dear Reviewer,

As we are approaching the end of the rebuttal period, we would like to cordially inquire about the extent to which we have successfully addressed the concerns outlined in your review.

Should there be any lingering points that require further attention, please rest assured that we are enthusiastic about the opportunity to provide comprehensive responses to any subsequent queries or comments you may have.

Your constructive input remains invaluable to us, and we appreciate your dedication to enhancing the quality of our manuscript. Thank you for your time and consideration.

Best,
Authors
>I am confused about the necessity of adopting a stack-based data structure to merge the local A_i and B_i modules. It seems that we can just adopt a set to store all modules.

While storing all modules in a set does indeed have the same memory cost as our proposed method, it remains impractical in real-world scenarios for two primary reasons:

**(1) Incompatibility with LoRA Implementation:** LoRA is not implementable when stored separately. In LoRA-fine-tuned language models, the LoRA modules can be utilized without merging them into the base model parameters, allowing the LoRA modules and the original model parameters to be stored together. As discussed in the LoRA paper [1], the forward pass is represented as:
                                                            
$h = \Delta x + W x=W x+BAx,$

where $BA$ represents the LoRA parameter. In our FLoRA approach, the A and B are global LoRA modules. However, if we store all modules separately, the inference stage would require the computation:

$h = \Delta x + Wx = Wx + B_0A_0x + B_1A_1x + B_2A_2x + … + B_{K-1}A_{K-1} x,$

This necessitates applying  $K$ LoRA modules within a single base model, which introduces iterative computations that are inefficient and impractical during the inference stage. Moreover, current LoRA codebases do not support this approach, where a single base model integrates multiple LoRA modules. The inconvenience of using a set to store LoRA modules becomes particularly problematic as the number of clients increases.
 
**(2) Incompatibility with Multi-Round Federated Fine-Tuning:** Storing LoRA modules to a set cannot effectively support multi-round federated fine-tuning. In subsequent rounds, fine-tuning requires an updated base model, necessitating updating the base model on the server or the clients. Updating on the server, as discussed in our rebuttal, incurs significantly higher communication costs. Conversely, updating on the clients imposes a substantial computational burden on each client, making this approach impractical.

>It appears that the communication cost is linear to the number of selected clients and the rank of A and B modules. Could the method still save communication costs when the number of selected clients is large or the rank is large?

Yes, we are indeed able to save communication costs when using a widely applied LoRA rank.

LoRA is commonly employed for fine-tuning models because it significantly reduces computational resource demands. This advantage hinges on the premise that the LoRA rank is much smaller than the rank of the model parameters themselves. The key prerequisite for using LoRA is that the number of parameters in LoRA must be substantially smaller than that of the base model. Otherwise, the benefits of LoRA would be negated, making full fine-tuning a more efficient option. Therefore, the LoRA rank should not be so large that it surpasses the communication resources required for transmitting the model’s original parameters.
 
Regarding the number of selected clients, it is important to note that the commonly used LoRA rank is typically about 1/100 of the base model size. For instance, in the case of Llama-7b, which has 4096x4096 full parameters, a typical LoRA configuration might use only 16x4096 parameters. This relationship can be expressed as:

$P_{LoRA} << P_{full}$

,where $P$ is the parameter size. When the clients send parameters to the server, FLoRA can consistently save communication costs because:

$P_{LoRA} < P_{full}$

This inequality holds true whenever LoRA is used. Moreover, when the server sends parameters back to the clients, FLoRA can also save communication costs under the condition:

$2K < P_{full}/P_{LoRA}$

Given that $P_{full}/P_{LoRA}$ is typically greater than 100 in fine-tuning scenarios, and federated fine-tuning server generally lack the communication resources to support extensive client communication (e.g., current works use 10 clients [2]), FLoRA proves effective in reducing communication cost.

References:

[1] Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., ... & Chen, W. (2021). Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685.

[2] Chen, J., Xu, W., Guo, S., Wang, J., Zhang, J., & Wang, H. (2022). Fedtune: A deep dive into efficient federated fine-tuning with pre-trained transformers. arXiv preprint arXiv:2211.08025.
Thanks for your response. I still have some concerns.

1. I am confused about the necessity of adopting a stack-based data structure to merge the local A_i and B_i modules. It seems that we can just adopt a set to store all modules. 

2. It appears that the communication cost is linear to the number of selected clients and the rank of A and B modules. Could the method still save communication costs when the number of selected clients is large or the rank is large?
Thank you for addressing the concerns raised during my review process, and I have also gone through all the comments from other reviewers. I agree with Reviewer 7eRV on the contribution of this work, the proposed method is simple yet effective. Our community indeed benefits from straightforward methods that address realistic problems. The simplicity of this work is not a drawback; rather, it facilitates broader use and adoption by the entire community.

The authors have successfully identified and rectified the mathematical errors present in current methods with their approach. Therefore, I believe this paper should be accepted by NeurIPS 2024 to attract broader attention from the community. More researchers need to explore the proposed method and potentially develop novel solutions to address the identified problem.

Based on the points mentioned above, I have decided to raise my score.
Thanks to the authors for the response and my concerns are well addressed. The simple yet effective method is friendly for practical applications, especially for the federated learning field, so I increase my score.