FedEx-LoRA: Exact Aggregation for Federated and Efficient Fine-Tuning of Foundation Models

Contribution/Missed Baselines

We appreciate the reviewer highlighting these works.

The paper "Improving LoRA in Privacy-Preserving Federated Learning" [1], referred to as FFA-LoRA in their work, was thoroughly addressed in our original manuscript. We explicitly discussed it in Section 3 (lines 170-176), included it as a baseline (lines 269-271), and provided detailed comparisons across all experiments. In the original manuscript, these comparisons appear in Tables 1, 2, 5, and 7, while in the new manuscript, they are found in Tables 1, 2, 3, 4, 9, 10, and 11. We are unsure how the reviewer missed these multiple mentions of FFA-LoRA in our original paper.

Our analysis covered its trade-offs between exactness and expressivity, and evaluated it on multiple settings. Our findings highlighted the impact of its choice to freeze the A matrices, as we consistently outperformed FFA-LoRA across all models and datasets.

Regarding FLORA [2], *per ICLR 2025 guidelines, papers published after July 1, 2024 are considered contemporaneous work. As FLORA was released in September 2024, while relevant, it falls outside the required comparison window for submissions. *Here is the link for ICLR's explicit policy: (https://iclr.cc/Conferences/2025/FAQ)

Experiments

We appreciate these suggestions and have addressed these points in our revision. We have expanded our work to evaluate larger models and more complex tasks. Our method significantly outperforms state-of-the-art federated LoRA approaches when evaluated on complex tasks with LLMs. This improvement is much more pronounced compared to our previous evaluations with smaller models. The manuscript now includes detailed evaluations of models ranging from 3B to 9B parameters, covering both arithmetic and commonsense reasoning tasks. For arithmetic reasoning, we fine-tuned Mistral-7B and Gemma-2 9B on a subset of the MetaMathQA dataset and evaluated their performance using the GSM8K and MATH benchmarks. For commonsense reasoning, we fine-tuned Llama-3.2 3B on COMMONSENSE170K—a collection of eight commonsense reasoning datasets—and reported performance for each task. The results, presented in Tables 1 and 2 (of the revised manuscript), demonstrate that our method consistently outperforms other methods across all tasks.

Method

Communication Cost

The communication overhead is carefully managed through matrix decomposition. Specifically, since ∆Wres has rank ≤ (kr), we can decompose it using Singular Value Decomposition (SVD) into U, S, V. Now the singular values will be 0 post row and column > (kr). We can simply take the first (kr) rows of U and similar columns for V. This decomposition means we only need to transmit matrices of size m×(kr) and (kr)×n, requiring (kr)(m+n) parameters. In contrast, transmitting the full weight update would require m×n parameters.

For concrete quantification, consider our RoBERTa-base experiments where m=n=768, k=3 clients, and r=4: FedEx-LoRA requires transmitting only 2×768×(3×4) = 18,432 parameters per layer, compared to 768×768 = 589,824 parameters for full model updates. This represents a 32x reduction in communication cost per layer (the total cost is proportional to the cost per layer). This analysis aligns with our empirical measurements in Table 4 (original manuscript) or Table 6 (revised manuscript), where FedEx-LoRA shows only marginal overhead compared to standard federated LoRA approaches while being significantly cheaper than federated full fine-tuning.

Computation Cost

The matrix multiplications required for aggregation are minimal compared to the computational demands of model training on clients. For perspective, a single forward pass by any client requires orders of magnitude more computation than the server's aggregation operations. The training computations dwarf the simple matrix operations needed for aggregation. Recent work in federated learning (Kairouz, et al.,2021) advocates for leveraging server computation to improve system efficiency. We believe that our minimal computational overhead should not be considered a weakness as it aligns well with modern FL approaches.

Writing

We appreciate the reviewer's suggestions to enhance the clarity of the manuscript. We have revised the equations in Section 4 for consistent notation. Superscripts now indicate the aggregation round, while subscripts denote the client number. Global weights and adapters are clearly specified in the subscript for clarity.

References:

[1] Sun, Y., Li, Z., Li, Y., & Ding, B. (2024). Improving LoRA in privacy-preserving federated learning. arXiv preprint arXiv:2403.12313

[2] Wang, Z., Shen, Z., He, Y., Sun, G., Wang, H., Lyu, L., & Li, A. (2024). Flora: Federated fine-tuning large language models with heterogeneous low-rank adaptations. arXiv preprint arXiv:2409.05976

[3] P. Kairouz, H. B. McMahan, B. Avent, et al., "Advances and open problems in federated learning," arXiv preprint arXiv:1912.04977, 2021

Thank you for your comments and feedback. We would like to specifically highlight the point regarding comparisons to FFA-LoRA and FLoRA. FFA-LoRA was thoroughly addressed in our manuscript through discussions, baselines, and detailed comparisons across multiple experiments, as explained in our rebuttal. Regarding FLoRA, as per ICLR 2025 guidelines, it falls outside the required comparison window, being released in September 2024.

We believe this clarification is critical, and we kindly request your response to this point before the discussion period ends. For the rest of your feedback, we hope our rebuttal and updated draft have addressed your questions and incorporated your suggestions. Please feel free to share any further questions or thoughts before the discussion period ends.

We thank the reviewer for the constructive feedback. We have carefully addressed each concern and question raised in the review.

Weaknesses

Novelty

While the method is straightforward, we believe this simplicity is a strength rather than a limitation. Like widely-adopted methods such as AdamW (Loshchilov, et al.,2019), FedEx-LoRA's value lies in its practical utility and ease of implementation while providing significant empirical improvements.

Scalability

We acknowledge the scaling concerns with larger client populations. We propose using truncated SVD, which maintains the same communication cost as FedIT while providing optimal approximation:

Best Inexact Approximation

For exact aggregation, the communication cost scales linearly with the number of clients, becoming prohibitive in hyperclient settings. To address this, we propose relaxing the exact aggregation condition through truncated SVD of the residual matrix. This reconstruction yields a low-rank approximation which, by the Eckart-Young theorem, is provably optimal for the high-rank update matrix. Specifically, for a target rank r', the best low-rank approximation ΔW_rec' is computed as:

U, S, V.T ← SVD(ΔW_res), ΔW_rec' ← U[1 : r'] S[1 : r', 1 : r'] V.T[1 : r']

While this method introduces approximation error, it provides the theoretically optimal approximation to exact aggregation. A key advantage is that the server can control communication costs, a capability absent in previous methods - FedIT and FFA-LoRA. The method has also been updated in the manuscript (Section 4, lines 249-261).

To test this efficient communication protocol, we evaluate our method in scenarios with a large number of clients where the original method would incur prohibitive communication costs.

We fine-tune Mistral-7B with 100 clients on a subset of MetaMathQA and evaluate on GSM8K and MATH benchmarks. Since communication cost scales linearly with the truncated SVD rank of the residual matrix and equals FedIT's cost at the same rank, our results are directly comparable in terms of communication overhead. The table below shows these results, where 'Rank' indicates the rank of the communicated matrix. Our method outperforms FedIT at equal communication bandwidth and maintains similar performance even when using half the bandwidth.

Questions

1. To assess performance across different settings, we expand our evaluation with additional experiments using 10 clients, randomly selecting 20% per round, and applying truncated SVD with rank 4 for aggregation. The results show promising scalability: FedEx-LoRA improves RoBERTa-base performance on CoLA from 54.17 to 55.26 and on MRPC from 85.21 to 86.27, outperforming FedIT. These findings suggest that our method maintains its advantages even with more clients, while keeping communication costs comparable to FedIT. Additionally, it achieves the theoretically optimal approximation, as evidenced by a reduction in the scaled Frobenius norm of deviation from 0.26 to 0.075 (this metric is clearly defined in Section 6 of the paper). Our experimental setup follows established protocols from recent federated LoRA works (Sun et al., 2024; Cho et al., 2024) to ensure fair comparisons. However, we agree that evaluating our method in diverse settings would help demonstrate its robustness across various scenarios.

2. The effect of the residual term can be indirectly measured or analyzed. Since the residual term added to the product of the aggregated matrices A and B is constant, modifying the residual term effectively changes the assignment of A and B after aggregation. In this context, we believe that the discussion in Section 6 on “Assignment Strategies for A and B” is exactly equivalent to analyzing the impact of varying the residual term.

[1] Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. arXiv preprint arXiv:1711.05101.

[2] Sun et al. "Improving LoRA in Privacy-preserving Federated Learning", arXiv:2403.12313, 2024

[3] Cho et al. "Heterogeneous LoRA for Federated Fine-tuning of On-Device Foundation Models", arXiv:2401.06432, 2024

We thank the reviewer for recognizing the strengths of our paper and providing valuable suggestions. We have updated the manuscript to address the concerns raised.

Model Limitations and Task Scope

How does FedEx-LoRA scale and perform with larger foundation models and complex tasks?

We appreciate these suggestions and have thoroughly addressed these points in our revision. We have expanded our work to evaluate larger models and more complex tasks. Our method significantly outperforms state-of-the-art federated LoRA approaches when evaluated on complex tasks with LLMs. This improvement is much more pronounced compared to our previous evaluations with smaller models. The manuscript now includes detailed evaluations of models ranging from 3B to 9B parameters, covering both arithmetic and commonsense reasoning tasks. For arithmetic reasoning, we fine-tuned Mistral-7B and Gemma-2 9B on a subset of the MetaMathQA dataset and evaluated their performance using the GSM8K and MATH benchmarks. For commonsense reasoning, we fine-tuned Llama-3.2 3B on COMMONSENSE170K—a collection of eight distinct commonsense reasoning datasets—and reported performance for each task. The results, presented in Tables 1 and 2 (of the revised manuscript), demonstrate that our method consistently outperforms state-of-the-art federated fine-tuning techniques across all arithmetic benchmarks and commonsense reasoning tasks. For example, FedEx-LoRA achieves an 8.64% improvement over FFA-LoRA and a 2.42% gain over FedIT in average accuracy on commonsense reasoning tasks.

Table 1: Results for Llama-3.2 3B on eight commonsense reasoning datasets, comparing various federated LoRA methods at rank r=32.

Table 2: Arithmetic reasoning performance on GSM8K and MATH for Mistral-7B and Gemma-2 9B, comparing various federated LoRA methods at rank r=32.

Question 1

What factors contribute to FedEx-LoRA outperforming centralized LoRA on certain benchmarks?

Training on smaller, distributed datasets through federated learning can sometimes prevent the model from overfitting, which may explain why it occasionally performs better than training on the complete centralized dataset.

Question 2

Can FedEx-LoRA be extended to other LoRA variants, such as MoE-LoRA and QLoRA, and if so, what adaptations might be necessary?

FedEx-LoRA integrates easily with QLoRA through a simple process. Each client first trains a model using QLoRA and then converts the trained weights to LoRA format using bfloat16 quantization. These converted weights are then sent to the server. The server aggregates the LoRA weights (including the residual matrix) and sends them back to the clients. Each client then loads the aggregated weights into their base model, converts the model back to QLoRA format, and continues training. The quantization and dequantization operations occur entirely on the client side. These operations are negligible compared to the costs of training a model, and they only take place before and after a federated aggregation round, not during the actual training phase.

MoE-LoRA can be used without modifications, requiring only that client models have the same number of LoRA experts. While the effectiveness of the aggregated model remains uncertain when LoRA experts may be trained on different tasks, we emphasize that FedEx-LoRA extends to any LoRA variant compatible with vanilla Federated LoRA, as our contribution lies in the aggregation mechanism.

Privacy Impacts

We thank the reviewer for pointing this out. This is indeed an interesting area of study, but we have focused our current work elsewhere and plan to explore this thoroughly in future studies.

We thank the reviewer for the constructive suggestions. We have addressed each point and concern raised in the review.

Communication Cost

The communication overhead is carefully managed through matrix decomposition. Specifically, since ∆Wres has rank ≤ (kr), we can decompose it using Singular Value Decomposition (SVD) into U, S, V. Now the singular values will be 0 post row and column > (kr). We can simply take the first (kr) rows of U and similar columns for V. This decomposition means we only need to transmit matrices of size m×(kr) and (kr)×n, requiring (kr)(m+n) parameters. In contrast, transmitting the full weight update would require m×n parameters.

For concrete quantification, consider our RoBERTa-base experiments where m=n=768, k=3 clients, and r=4: FedEx-LoRA requires transmitting only 2×768×(3×4) = 18,432 parameters per layer, compared to 768×768 = 589,824 parameters for full model updates. This represents a 32x reduction in communication cost per layer (the total cost is proportional to the cost per layer). This analysis aligns with our empirical measurements in Table 4 (original manuscript) or Table 6 (revised manuscript), where FedEx-LoRA shows only marginal overhead compared to standard federated LoRA approaches while being significantly cheaper than federated full fine-tuning.

Computation Cost

The matrix multiplications required for aggregation are minimal compared to the computational demands of model training on clients. For perspective, a single forward pass by any client requires orders of magnitude more computation than the server's aggregation operations. The training computations dwarf the simple matrix operations needed for aggregation. Recent work in federated learning (Kairouz, et al.,2021) advocates for leveraging server computation to improve system efficiency. We believe that our minimal computational overhead should not be considered a weakness as it aligns well with modern FL approaches.

References:

[1] Li, Z., Yan, C., Zhang, X., Gharibi, G., Yin, Z., Jiang, X., & Malin, B. A. (2023). Split Learning for Distributed Collaborative Training of Deep Learning Models in Health Informatics. arXiv preprint arXiv:2308.11027.

[2] P. Kairouz, H. B. McMahan, B. Avent, et al., "Advances and open problems in federated learning," arXiv preprint arXiv:1912.04977, 2021.

Thank you for your comments and questions. We hope the rebuttal adequately addresses your feedback. Additionally, we’ve included new results in the revised version and would appreciate your thoughts or further discussion on this before the discussion period concludes.