Robust Federated Finetuning of LLMs via Alternating Optimization of LoRA

Thanks for the detailed review. We address your concerns below:

Federated Round Definition, Communication and Computation Efficiency. Thanks for the helpful comments. In our paper, each communication round refers to one upload/download of either matrix A or B. RoLoRA updates and transmits only one matrix per round, halving both communication and computation compared to FedIT—see table.

We apologize for any confusion from the use of "iteration" in Algorithm 1, which was for alignment with Algorithm 2 and theoretical clarity. All experiments were rigorosly benchmarked using communication rounds. We will revise the manuscript to clearly distinguish rounds from iterations and add a footnote in Algorithm 1 to clarify this.

Rank-1 Analysis and Sample Complexity. Thanks for the insightful questions.

Regarding sample complexity, Theorem 5.4 shows that $q$ scales with $d$ for rank-1, and generalizes to $O(dr)$ for higher ranks. This follows from the $\epsilon$-net argument (Eq. 34), where the covering number grows exponentially with dimension $dr$ [1, Sec. IV.A], leading to the updated $q$ (Line 755). The sample complexity can be justified that in models like $Y=Xab^\top$, $2d$ unknown parameters necessitate sample complexity proportional to $d$.

The rank-1 case highlights FFA-LoRA's core limitation—its inability to align down-projection—which persists empirically in higher ranks (see Sec. 5.2), in language tasks. We further discuss the generalization with Reviewer JAAe under Rank-1 Limitation and Generalizability, and note the core issue of constrained expressivity remains unchanged in the higher-rank case.

[1]Nayer, S., & Vaswani, N. (2022). Fast and sample-efficient federated low rank matrix recovery from column-wise linear and quadratic projections. IEEE Trans. Inf. Theory

Higher variance in Fig.3. Thanks for the question. The higher variance in Fig. 3 before convergence is expected due to different initializations, which can lead to varying optimization trajectories. After convergence, RoLoRA shows low variance, consistent with its stable performance in Tables 1 and 2. FFA-LoRA's variance differs across Tables 1 and 2 mainly due to task difficulty—it's low on simpler tasks (e.g., GLUE, BoolQ) but higher on complex ones (e.g., PIQA, SIQA) due to sensitivity to initialization, aligning with its weaker performance in Table 1 and 2.

Convergence curve for 100 Communication rounds. The figure shows the 100-round extension of Fig. 3, where RoLoRA consistently converges faster and achieves the highest accuracy.

Fig. 3 aimed to compare convergence under a fixed sample budget, not full convergence. As Table 2 shows, this budget is sufficient for all methods with 3 clients, but only RoLoRA fully converges with 50 clients—highlighting its efficiency in low-resource settings.

Differences from FLoRA and FedIT. Thanks for the questions. FedIT is equivalent to our FedAVG with LoRA baseline ("LoRA" in our paper), and we provide extensive comparisons. We'll explicitly cite and discuss FedIT in the revision.

We already discuss FLoRA (Lines 31, 113, 162). FLoRA shares RoLoRA's motivation but differs in approach: FLoRA aggregates full matrix products of LoRA-A and LoRA-B, while RoLoRA freezes one matrix for efficient, exact updates (Eq. 3-4). See Sec. 3.2 for discussion.

We've added table comparing RoLoRA and FLoRA under IID setting. In the 3-client setting, we ran 500 rounds and scaled rounds down proportionally with more clients to keep the total sample budget fixed. RoLoRA consistently outperforms FLoRA across tasks and client counts. While FLoRA eventually converges (e.g., 83.3% on MNLI after 4000 rounds), it does so much more slowly, highlighting RoLoRA's faster convergence and better scalability.

Please also see table comparing the communication cost and time cost. RoLoRA and FFA-LoRA have the lowest communication/time costs, while FLoRA is much more expensive.

Fig.1 Improvement and Use of Large in RoBERTa-Large. Thanks for the suggestions. We've updated the figure and will refer to RoBERTa-Large simply as a language model, without implying scale beyond its name.

Related Works. We agree and will revise the related work section to clearly highlight RoLoRA's key differences from prior methods.

RoLoRA on Linear Model. As suggested, we ran an experiment removing the ReLU from NN on MNIST—see figure. Across both linear and non-linear settings, all methods perform similarly, with RoLoRA showing modest improvement in the non-linear case, likely due to its better utilization of the added expressiveness from ReLU.

Thanks for the thoughtful review. We appreciate the recognition of RoLoRA's contributions and the strengths of our alternating optimization strategy, theoretical insights, and empirical results. We address the main concerns below:

Limited Theoretical Scope. We appreciate the reviewer's thoughtful feedback regarding the scope of our theoretical analysis. First, we would like to clarify that our convergence analysis, including comparisons between RoLoRA and FFA-LoRA, also covers the heterogeneous setting (Appendix A.5 at Line 1430). The analysis mirrors the homogeneous case: RoLoRA reduces the angle to the ground truth $\mathbf a$, ensuring convergence, while FFA-LoRA's loss remains tied to its initial angle. This has been discussed at Line 281 in maintext.

While our theoretical analysis focuses on the rank-1 LoRA, linear regression setting, this case remains foundational and highly non-trivial. By directly comparing the solutions of RoLoRA and FFA-LoRA under this simplified yet fundamental scenario, we rigorously demonstrate the inherent limitations of FFA-LoRA in representing low-rank updates. This provides critical insights into the expressiveness of FFA-LoRA even in its most basic form, establishing a baseline for understanding its behavior. Our empirical validation on neural network and MNIST shows the theorem can be extended to higher-rank and non-linear setting. Critically, the results on realistic LLM tasks bridges theory and practice, showing that the phenomena observed in our theoretical framework persist in non-linear, non-convex settings. This alignment mirrors methodologies in prior works [1, 2], which adopt simplified setups to distill core principles before validating them in real-world contexts.

A full theoretical comparison of federated LLM algorithms across all settings is currently infeasible. For neural networks, while convergence analyses exist, they rely on loss landscape assumptions (e.g., smoothness, PL-conditions) that preclude direct comparison of optima between algorithms. Instead, our work focuses on a tractable case where direct comparisons are feasible, allowing us to uncover provable limitations that also hold in real-world scenarios.

[1]Collins, Liam, et al. "Exploiting shared representations for personalized federated learning." (ICML2021).

[2]Collins, Liam, et al. "Fedavg with fine tuning: Local updates lead to representation learning." (NeurIPS2022)

Assumptions on decreasing angles. Thank you for pointing this out. To clarify, although we assume a decreasing sequence of angles in Lemma 5.3, this result is used as an intermediate step in the proof of Theorem 5.4. In particular, we demonstrate in the proof of Theorem 5.4 at Line 1027 that the angle decreases in the first iteration. Building on this, we apply an inductive hypothesis to show that the decreasing trend holds for all subsequent iterations (see Line 1039). In summary, the decreasing angle sequence is not assumed in the proof of the main result (Theorem 5.4), but is instead derived as part of the argument.

Algorithm used for theoretical analysis. Lemma 5.3 is tied to the update rule in Alg. 2. We will revise the lemma's description to make this connection more explicit.

Error bound on a. Thank you for the insightful question. The error bounds in Lemma 5.3 and Theorem 5.4 are derived based on Alg. 2, which updates both $\mathbf{a}$ and $\mathbf{b}$. The reason the bound applies specifically to the angle distance of $\mathbf{a}$ is that $\mathbf{b}$ is fully optimized at each client, while $\mathbf{a}$ is updated via gradient descent. This ensures that $\mathbf{b}$ is always optimal with respect to the current $\mathbf{a}$, allowing us to focus the analysis on the convergence behavior of $\mathbf{a}$. We designed RoLoRA for the linear regressor using alternating minimization (for $\mathbf{b}$) and gradient descent (for $\mathbf{a}$) to decouple their updates and eliminate the potential impact of insufficient local updates of $\mathbf{b}$ on the convergence of $\mathbf{a}$. We will clarify this motivation and its implications in the revised manuscript.

We hope these clarifications address the reviewer's concerns. Please feel free to reach out if the reviewer has any further concerns, and we would be glad to discuss them with the reviewer.

We thank the reviewer for their thoughtful and detailed assessment of our paper. We are encouraged by the overall positive evaluation and would like to clarify several points.

Rank-1 Limitation and Generalizability to higher ranks. Thanks for pointing this out. Although our analysis is conducted under the rank-1 setting, it remains highly nontrivial. In particular, we provide a direct and rigorous comparison between the solutions of RoLoRA and FFA-LoRA, clearly demonstrating the limitations of FFA-LoRA in this fundamental case. This already yields valuable theoretical insights into the core differences between the two methods. Prior works on similar algorithms and settings (See Sec 4.1 and 4.2 in [1], and Appendix A.3.1 and A.3.2 in [2]) have successfully analyzed both rank-1 and higher-rank cases using comparable techniques. This suggests that the rank-1 and higher-rank analyses share underlying structures and intuition. We believe our proof can be extended to higher ranks with additional technical work, but doing so would not change our main conclusion regarding FFA-LoRA's limited expressiveness. We outline key steps toward this extension and leave it for future work.

Orthonormalization in Algorithm. For the rank-1 case, it is only required to normalize the updated $a$ to unit length (Line 12 of Algorithm 2). To maintain orthonormality for the higher-rank case, we need to include a QR step $A^+=AR$ where $A$ is updated from through GD (cf. Line 11 of Algorithm 2), ${A^+}^\top A^+ = I_r$, and $R$ is upper-triangular.

Error metric. For the rank-r case, we define the subspace distance of two $d\times r$ matrices (with orthonormal columns) as the following:${SD}( A, A^*) = ||( I_ d- A A^\top ) A^*||_ F$ This is a direct generalization of the rank-1 case. Geometrically speaking, ${\rm SD}( A, A^*) = \sqrt{\sum_{r'=1}^r \sin^2(\theta_i)}$, where $\theta_1,\cdots, \theta_r$ are the principal angles between the column spaces of $A$ and $A^*$.

Generalizing Eq. (17). For rank-$r$, we similarly bound $|\bar{{B}} - {G}|_{op}$ by ${\rm SD}({A}, {A}^*)$, where ${G} = B(A^*)^\top {A}$. The main technical step is controlling $|({A}^\top {X}_i^\top {X}_i {A})^{-1}|{op}$, handled via Bernstein's inequality and an $\epsilon$-net (see [3], Sec. IV-E).

Condition-Number Dependecies. In the rank-1 case, the convergence and sample complexity depend on the norm of the ground-truth signal vector $b^*$, as shown in Eq. (13,44,123). For higher-rank settings, this dependency generalizes to the operator norm (i.e., the largest singular value). However, the operator norm alone does not fully characterize the problem complexity. Instead, the condition number, which is the ratio of largest and smallest singualr value, is critical. The condition number is reduced to 1 when rank=1.

[1]Jain et al., "Low-rank matrix completion using alternating minimization," STOC 2013

[2]Thekumparampil et al., "Statistically and computationally efficient linear meta-representation learning," NeurIPS 2021

[3]Nayer, S., & Vaswani, N. (2022). Fast and sample-efficient federated low rank matrix recovery from column-wise linear and quadratic projections. IEEE Trans. Inf. Theory

Assumptions on decreasing angles. Thank you for pointing this out. While Lemma 5.3 assumes a decreasing angle sequence, it's only an intermediate step for Theorem 5.4. In Theorem 5.4's proof (Line 1027), we show the angle decreases in the first step and use induction to extend this to all iterations (Line 1139). Stochastic setting: Our analysis is deterministic; extending to the stochastic case is left for future work. We'll clarify this in the revision.

Difference between Lemma A.11 and the first part of the proof of Theorem 5.4. Thank you for the careful reading. Eq. (123) and Eq. (166) are essentially the same—Eq. (123) assumes a decreasing angle, while Eq. (166) proves this for the first iteration. Setting $t=0$ in Eq. (123) gives Eq. (166). We restated the proof for clarity and completeness but will revise it to be more concise.

Eq. (87). Thank you for pointing this out. Eq. (87) follows by normalizing Eq. (86). The reference to Eq. (87) on Line 929 was a typo—it should be Eq. (88). We'll correct this in the revision.

Performance in Table 12. Thank you for the observation. LLaMA-2-7B is already strong on MMLU due to its pretraining. Since MMLU focuses on factual recall, not task-specific adaptation, PEFT methods like LoRA, RoLoRA, or FFA-LoRA offer limited gains. This is also observed at Sec 5.3 in [4]. In contrast, we show clear improvements on task-specific adaptation tasks such as GLUE and commonsense reasoning tasks.

[4]Guo, Pengxin, et al. "Selective Aggregation for Low-Rank Adaptation in Federated Learning." (ICLR 2025).

Related works. Thanks for the helpful suggestions and for taking the time to check for related work. We will cite them in the related work section.

Thank you for the detailed and constructive review. We address the key concerns below:

FFA-LoRA's Reduced Expressiveness and RoLoRA's Theoretical Superiority. Thank you for the constructive comments. We show in Proposition 5.5 (proof in Appendix A.4.1) that FFA-LoRA is suboptimal in our linear setting. This result holds for any unit vector $\mathbf{a}$ and corresponding $\mathbf{b}$ obtained by fully minimizing the local loss, as in Line 5 and 7 of Algorithm 2. The same expected loss bound applies to RoLoRA. As discussed in Line 271-280, substituting RoLoRA's reduced angle $\epsilon$ at Line 259 into Eq. (11) shows its expected loss can be made arbitrarily small, unlike FFA-LoRA, whose loss is limited by the initial angle.

Scope of Convergence Analysis Our convergence analysis also extends to the heterogeneous setting (Appendix A.5), following the same logic as the homogeneous case. RoLoRA reduces the angle between $\mathbf{a}$ and $\mathbf{a}^*$, ensuring convergence to the global optimum, while FFA-LoRA's loss remains limited by its initial angle, as discussed at Line 281.

Though based on a simplified linear model, the proof is non-trivial. The linear setting allows direct comparison due to its unique global minimum, unlike neural networks, where only convergence to local minima can be shown and final losses across methods are not directly comparable.

Non-IID robustness. Thank you for the constructive comments. We provide a theoretical analysis under a heterogeneous linear setting (Appendix A.5) and evaluate RoLoRA's robustness to non-IID data using a two-layer neural network (Fig. 2). Additionally, we ran experiments on a language model under non-IID conditions—see the table. RoLoRA consistently outperforms LoRA, FFA-LoRA, and FlexLoRA on MNLI and QQP across varying data heterogeneity and client counts.

Re-Organization of Proof. Thank you for the thoughtful feedback. We already provide a high-level overview of the main proof at Line 271. To improve clarity, we will revise the appendix to include a more detailed outline and highlight key intuitions behind the technical steps.

Differential Privacy. Thanks for the suggestion. We itegrated NbAFL [1] and the results are in the table. We use $\epsilon = 10, \delta = 1e-6$. In this setting, RoLoRA outperform others across MNLI and QQP tasks.

[1]Wei, Kang, et al. "Federated learning with differential privacy: Algorithms and performance analysis." IEEE transactions on information forensics and security 15 (2020)

Extreme low-rank stability. Thank you for raising this point. Our experiments already include rank 1, 2, 4, and 8 settings (see Fig. 4, 6-8 in Appendix B.2.2 at Line 1743), demonstrating RoLoRA's robustness under tight parameter budgets. Table 1 and 6 show results with increasing clients at rank 4 and 8. We've also added experiments with rank 2 and different numbers of clients—see the table and figure—where RoLoRA still shows strong convergence and competitive accuracy.

Ablation study of the role of A. Thank you for the insightful suggestion. To address this, we conducted an experiment comparing performance of FFA-LoRA, RoLoRA, and different mixing strategies under the setting with 50 clients. In these strategies, for example, 20%RoLoRA+80%FFA-LoRA means we finetune with RoLoRA (where A is learned) for the first 20% of communication rounds, followed by FFA-LoRA (where A is frozen) for the remaining 80%. The results are shown in the figure. We observe that finetuning with RoLoRA generally leads to faster convergence and higher final accuracy, highlighting the benefits of learning A, especially in early training.

Comparison with FlexLoRA. Thank you for your comment. Our paper already compares RoLoRA and FlexLoRA across ranks and client counts (Table 6, Line 1705). We've added convergence results in a 50-client setting figure and a table comparing communication and time costs. Results show RoLoRA's clear advantages in large-scale, low-resource settings.

Related works. Thanks for taking the time to check for related work. We will discuss Koo et al. (2024) in the related work section.