Riemannian Low-Rank Adaptation for Federated Fine-Tuning of Foundation Models

Table 2: Performance comparison of methods on ViT + CIFAR-10

Weaknesses 2: Key parameters, like rank thresholds and learning rates, are not thoroughly analyzed. This could afect the method’s adaptability to various tasks and models.

Answer: Thank you for pointing this out. In the submission, we have already conducted a detailed analysis of rank thresholds and learning rates, as shown in Tables 26-29 in Appendix 7.6.

Tables 26-27 in Appendix 7.6 evaluate the performance effects by varying the rank selection threshold $\Theta$ from 0.6 to 1. It is observed that the performance score initially increases quickly and then become stable or even drop when $\Theta$ continuously increases. Initially, a large $\Theta$ can help generate a large rank and utilize the strength of more model parameters for ensuring better learning results. Later on, when $\Theta$ continues to increase and goes beyond some thresholds, the performance score becomes stable or drop. A rational guess is that excessive model parameters may bring redundant information or even noise, leading to performance degradation. It is important to find the optimal rank selection threshold to guarantee better model performance.

Tables 28-29 in Appendix 7.6 evaluate the performance effects by changing the learning rate ($LR$) from 0.0001 to 0.01. We have observed that for Roberta-based tasks, higher learning rates generally lead to better performance. SST-2, MRPC, and MPQA achieve the highest accuracy at $LR$=0.005 or 0.01. For LLaMA-based tasks, moderate learning rates, such as 0.001 or 0.0005, can produce the optimal results, while lower learning rates (e.g., 0.0001 or 0.0005) lead to better learning results for ViT-based tasks. This demonstrates that it is crucial to choose the optimal learning rates for the training to achieve the competitive performance.

We thank this reviewer for the great suggestion!

Weaknesses 1: The efficiency of Riemannian operations is not fully quantified. A clearer breakdown of resource usage would help in assessing practical overhead.

Answer: Thank you for your suggestion. By following existing work Beyond Efficiency: A Systematic Survey of Resource-Efficient Large Language Models, we conduct the experiments to evaluate the resource usage of our RAFFT method across various metrics, including FLOPS (measure computational complexity), throughput (assess the processing rate), communication volume (quantify data transfer during training), training time, and accuracy. We have included these results in Appendix 7.6 table 34-35.

As shown in the following table 1 and 2, the resource usage of 16 federated parameter-efficient fine-tuning methods is evaluated on the LLaMA 7B + MPQA task. Our RAFFT method achieves the highest throughput of 6,229.19 tokens/sec and the comparable FLOPS with the best baselines (72.6401 GFLOPS). In addition, RAFFT significantly outperforms all other LoRA-based methods with the lowest communication volume of 3,146,112 bytes, and achieves the highest accuracy of 90.8%. Similarly, on the ViT + CIFAR-10 task, RAFFT achieves the best accuracy of 61.22%, with the shortest training time of 3,222 seconds and a communication volume of 98,352 bytes. These results validate our RAFFT method is able to well balance computational and communication efficiency. A plausible explanation for this superior performance is that the Riemannian parameter matching and Riemannian gradient optimization enable faster convergence and improved accuracy, while the adaptive rank selection reduces redundant computation, making RAFFT a highly practical and efficient approach for real-world federated learning scenarios.

Table 1: Performance comparison of methods on LLaMA 7B + MPQA

Fourth, Eq.(59) gives the orthogonal projection of any update direction $Z^k$ onto the tangent space $T_{M_r}$ at $\Delta W^k$. Lemma 7 guarantees that the projection respects the manifold structure and maintains the rank $r$. The resulting Riemannian gradient in Eq.(65) incorporates these projections to guide updates within the manifold.

Finally, to construct the retraction function that maps the tangent space updates back onto the manifold, Eq.(68) minimizes the distance between the perturbed point and the manifold. During this process, Lemma 8 uses the truncated singular value decomposition (SVD) to ensure that the rank is preserved throughout the projection. The QR factorizations of the augmented matrices $[U^k, U_p]$ and $[V^k, V_p]$ in Eq.(70), yield the orthogonal matrices $Q_U$, $Q_V$ and upper triangular matrices $R_U$, $R_V$, which facilitate the computation of the perturbed matrix $\Delta W^k + H$ in Eq.(71). Applying SVD to this perturbed matrix produces the retracted point $f(H)$, as described in Eq.(72). Eq.(73) provides the updated components $(U^{k+1}, \Sigma^{k+1}, V^{k+1})$ in the RGD optimization.

In summary, Theorem 2 rigorously proves the RGD optimization process on the Riemannian manifold $M_r$, ensuring rank invariance throughout the local update process while seamlessly integrating gradient computation, projection, and retraction operations.

Lemma 9 establishes a key inequality that relates the geodesic distance between successive points on a Riemannian manifold $\mathcal{M}$ to the reduction in the objective function during gradient-based optimization. It captures the relationship between the gradient $g^{(t)}$, the geodesic distance $\text{dist}(x, x^{(t)})$, and the step size $\alpha$, while accounting for the manifold's curvature through the factor $\zeta(\kappa_{\min}, \text{dist}(x, x^{(t)}))$.

Lemma 9 (Geodesic Distance Properties): For any Riemannian manifold $\mathcal{M}$ where the sectional curvature is lower bounded by $\kappa_{\min}$ and any point $x, x^{(t)} \in \mathcal{M}$, the update $x^{(t+1)} = \text{Exp}\_{x^{(t)}}(-\alpha g^{(t)})$ with $g^{(t)} \in T_{x^{(t)}} \mathcal{M}$ satisfies

$

\left\langle -g^{(t)}, \text{Exp}_{x^{(t)}}^{-1}(x) \right\rangle \leq \frac{1}{2 \alpha} \left(\text{dist}^2(x, x^{(t)}) - \text{dist}^2(x, x^{(t+1)}) \right) + \frac{\zeta(\kappa_{\min}, \text{dist}(x, x^{(t)})) \alpha}{2} |g^{(t)}|^2

$

Where $\text{Exp}\_{x^{(t)}}(x)$ is the exponentail map that projects a tangent vctor x onto the manifold, $\text{Exp}_{x^{(t)}}^{-1}(x)$ is the inverse exponential map that maps a point x back to the tangent space, $\alpha$ is the learning rate controlling the magnitude of the update, $\text{dist}(x, x^{(t+1)})$ denotes the geodesic distance between points x and $x^{(t+1)}$ on the manifold, and $\zeta(\kappa, c) = \frac{\sqrt{|\kappa|c}}{\tanh\left( \sqrt{|\kappa|c} \right)}$ is the curvature adjustment for distance.

Please refer to [6] for the detailed proof.

Theorems 3 and 4 conduct the convergence analysis of our RAFFT algorithm based on the RGD optimization on the Riemannian manifold in the settings of both non-convex and geodesically convex.

For the non-convex case, Definition 2 introduces the smoothness property of the loss function $L$ on a manifold, ensuring bounded gradients in Eq.(74). Theorem 3 establishes the convergence of the optimization process, demonstrating that the expected squared norm of the Riemannian gradient decreases to a bounded region (Equation 79). This result is achieved by projecting updated points back to the manifold using the retraction operator and measuring distances via the geodesic inverse in Eqs.(80)-(82). By leveraging the Lipschitz smoothness property (Equation 75), the inequality in Equation 83 bounds the reduction in the loss function $L$ over iterations. Summing these inequalities and selecting appropriate step sizes ensures convergence, as shown in Eq.(85).

For the convex scenario, Definition 3 ensures the function $L$ satisfies a convex-like property along geodesics, enabling convergence guarantees. For geodesically convex functions, Theorem 4 extends the analysis by proving the convergence rate for the difference between the current and optimal loss values. Lemma 9 is used to bound the geodesic distance between successive points and relates it to the reduction in loss in Eqs.(86)-(97). The result concludes with a convergence guarantee dependent on the curvature constant $\zeta$ and the geodesic distance to the optimal solution.

Lemma 5 ensures that the tangent space $T_M$ is well-defined and $M_r$ is an embedded submanifold of $\mathbb{R}^{m \times n}$, and thus introduces Definition 1 in the submission.

Lemma 5 (Differentiability of Local Defining Function): Let $h: \mathbb{R}^{(m-r) \times (n-r)} \to \mathbb{R}^{m \times n}$ be defined as: $h(Y) = Y_{22} - Y_{21} Y_{11}^{-1} Y_{12}$ where $Y_{ij}$ are submatrices of $Y$. Then $h$ is smooth and has an inverse mapping $h^{-1}$ such that $\text{ker}(Dh(Y))$ is a linear subspace that spans $\mathbb{R}^{m \times n}$.

Please refer to [4] for detailed proof.

Lemma 6 characterizes the dynamics of the corresponding smooth curves $\mathbf{U}(t)$ and $\mathbf{V}(t)$ of any vectors within $T_\mathbf{U}St(m, r)$ and $T_\mathbf{V}St(n, r)$ respectively, when they evolve smoothly on their respective manifolds.

Lemma 6 (Tangent Space on the Stiefel Manifold): Let $\mathbf{U}^k \in \mathbb{R}^{m \times r}$ and $\mathbf{V}^k \in \mathbb{R}^{n \times r}$ lie on the Stiefel manifolds $\text{St}(m, r)$ and $\text{St}(n, r)$, respectively. For any $\Omega \in \text{Skew}(r)$ and $B \in \mathbb{R}^{(m-n) \times r}$, the tangent space velocities $\mathbf{U}^{k\prime}(0)$ and $\mathbf{V}^{k\prime}(0)$ are given by:

$

\mathbf{U}^{k\prime}(0) = \mathbf{U}^k \Omega + \mathbf{U}^k_\perp B, \quad \mathbf{V}^{k\prime}(0) = \mathbf{V}^k \Omega' + \mathbf{V}^k_\perp C,

$

where $\mathbf{U}^k_\perp$ and $\mathbf{V}^k_\perp$ are orthogonal complements.

Please refer to [5] for the detailed proof.

Lemma 7 decomposes $Z^k$ into components within the tangent space, ensuring the projection results in a matrix that respects the manifold structure of a manifold $M_r$. It is utilized to compute the gradients in our proposed Riemannian gradient descent algorithm.

Lemma 7 (Orthogonal Projection onto Tangent Space): Let $\Delta W^k = U^k \Sigma^k (V^k)^\top$ represent a point on $M_r$. The orthogonal projection of a matrix $Z^k \in \mathbb{R}^{m \times n}$ onto the tangent space $T_{M_r}$ at $\Delta W^k$ is given by:

$

\text{Proj}_{\Delta W^k}(Z^k) = U^k G (V^k)^\top + U^k_\perp (V^k_\perp)^\top + U^k (V^k_\perp)^\top,

$

where $G$ is a general matrix relaed to the variation of $U^k$ and $V^k$ along their tangent directions, $U^k_\perp$ and $V^k_\perp$ denote the components of $Z^k$ orthogonal to $U^k$ and $V^k$, respectively.

Please refer to [5] for detailed proof.

Lemma 8 provides the theoretical foundation for projecting the updated matrix back onto the manifold after applying a gradient step in Eq.(69). The Eckart-Young-Mirsky theorem ensures that the retraction operation preserves the rank structure.

Lemma 8 (Retraction for Fixed-Rank Matrices): Let $H \in \mathbb{R}^{m \times n}$ represent the perturbation in the ambient space. The retraction function $f$ maps the tangent space $T_M$ to the manifold $M$ by minimizing the Frobenius norm distance:

$

f(H) = \arg \min_{\Delta W \in M} | \Delta W^k + H - \Delta W |_F^2.

$

Under the constraint of fixed rank $r$, the retraction is given by the truncated singular value decomposition of $\Delta W^k + H$.

Please refer to [1] for the detailed proof.

Theorem 2 illustrates that the RGD optimization on the Riemannian manifold ensures the rank invariance during the local update process.

First, Eq.(43) defines a manifold $M_r$ as an embedded submanifold of $\mathbb{R}^{m \times n}$, where the rank constraint $r$ is preserved through a matrix partitioning approach. Lemma 4 is leveraged to decompose the parameter matrices $W$ with rank $r$ into four submatrices $w_{11}, w_{12}, w_{21}$ and $w_{22}$, while these submatrices still maintain the invariance of $r$, as shown in Eq.(46).

Second, Eq.(47) defines a local parametrization function $h$ to formalize the manifold structure and ensure smooth updates. Its differentiability and inverse properties, supported by Lemma 5, guarantee that the manifold $M_r$ is smooth and that local updates remain within the manifold.

Third, during the RGD optimization process, tangent vectors are characterized through the evolution of $U^k$ and $V^k$ on the Stiefel manifolds $\text{St}(m, r)$ and $\text{St}(n, r)$. Lemma 6 characterizes the smooth dynamics of $\mathbf{U}^{k\prime}$ and $\mathbf{U}^{k\prime}$ in Eq.(53), ensuring orthogonality and the rank-preserving properties of these updates.

Weaknesses 4: The paper lacks necessary lemmas that would aid in understanding the proof sketches.

Answer: Thanks for the kind suggestion. We will include the following lemmas in the submission for better understanding of the theorem proof. Here, we will introduce and explain the lemmas one by one.

Lemmas 1-3 are the preliminary steps of proof of Theorem 1. Lemmas 1-3 and Theorem 1 together derive the upper bound on the Frobenius norm error for off-diagonal elements after an orthogonal transformation in Eq.(14) in the submission. We have included all of the proof of these lemmas and theorem in Appendix 7.2 ``Approximation Matrix Upper Bound".

Lemma 1 demonstrates the Frobenius norm of a matrix orthogonally invariant. It is used to conduct the approximation error analysis in Eq.(30).

Lemma 1 (Orthogonal Matrix and Frobenius Norm Preservation): Let $U \in \mathbb{R}^{n \times n}$ be an orthogonal matrix, i.e., $U^\top U = I$. For any $A \in \mathbb{R}^{n \times m}$, the Frobenius norm $\|A\|_F$ satisfies: $\|U^\top A\|_F = \|A\|_F.$

Please refer to [1] for the detailed proof.

Lemma 2 is the extension of the Cauchy-Schwarz inequality on matrices. It is utilized to derive an upper bound of Eq.(31) and further the bounded error of matrix approximation.

Lemma 2 (Cauchy-Schwarz Inequality for Frobenius Norm): For any $A, B \in \mathbb{R}^{n \times m}$, the Frobenius norms $\|A\|_F$ and $\|B\|_F$ satisfy: $|\text{tr}(A^\top B)| \leq \|A\|_F \|B\|_F.$ where $\text{tr}(\cdot)$ is the trace of a matrix.

Please refer to [2] for the detailed proof.

Lemma 3 validates how the condition number of $S^k$ governs the amplification of off-diagonal errors during inversion, thereby quantifying the impact on the approximation error. It is employed to obtain the total approximation error bound in Eq. (35) in Theorem 1.

Lemma 3 (Condition Number and Off-Diagonal Error Amplification): Let $S \in \mathbb{R}^{n \times n}$ be a symmetric positive definite matrix with condition number $\kappa(S) = \lambda_{\max}(S) / \lambda_{\min}(S)$. Then the Frobenius norm of the off-diagonal part of the inverse matrix is bounded by: $\|S^{-1}\_{\text{off-diagonal}}\|_F \leq \kappa(S) \|S\_{\text{off-diagonal}}\|_F.$ where $S^{-1}$ is the inverse of the matrix $S$.

Please refer to [3] for the detailed proof.

Theorem 1 obtains the upper bound on the Frobenius norm error for off-diagonal elements after an orthogonal transformation in Eq.(14) in the submission. Eq.(14) consists of three terms: the aligned parameter matrices $\mathbf{U}^k \mathbf{R}^k$ and $\mathbf{V}^k \mathbf{S}^k$ for client $k$ and the modified version $\tilde{\mathbf{\Sigma}}^k$ of diagonal matrix $\mathbf{\Sigma}^k$. The approximation error in Eq.(15) in Theorem 1 is raised by all the above three terms together. Eq.(29) leverages Lemma 1 to separate three components from the left side in Eq.(15). Eqs.(30)-(33) utilize Lemma 2 to estimate the error bounds led by $\mathbf{U}^k \mathbf{R}^k$ and $\mathbf{V}^k \mathbf{S}^k$. Eq.(35) employs Lemma 3 to derive the amplification of off-diagonal errors by $\|\tilde{\mathbf{\Sigma}}^k - \mathbf{\Sigma}^k \|_F$.

Lemmas 4-8 are the preliminary steps of proof of Theorem 2. Lemmas 4-8 and Theorem 2 together demonstrates that the Riemannian gradient descent (RGD) optimization on the Riemannian manifold ensures the rank invariance during the local update process. We have included all of the proof of these lemmas and theorem in Appendix 7.3 ``Riemann Gradient and Retraction".

Lemma 4 analyzes the correlation between a submatrix $\mathbf{w}_{11}$ of a matrix $\mathbf{W}$ and other submatrices of $\mathbf{W}$, which is essential for constructing the local defining function for the smooth manifold $M_r$. Therefore, it underpins the theoretical foundation of the matrix decomposition in Eq.(46).

Lemma 4 (Partitioning and Submatrix Inversion): Let $W \in \mathbb{R}^{m \times n}$ with rank $r$. If $w_{11} \in \mathbb{R}^{r \times r}$ is an invertible submatrix of $W$, then $W$ can be partitioned as:

$

W = \begin{bmatrix} w_{11} & w_{12} \\\ w_{21} & w_{22} \end{bmatrix},

$

where $w_{12}$, $w_{21}$, and $w_{22}$ are submatrices. The submatrix $X \in \mathbb{R}^{r \times (n - r)}$ satisfies:

$

X = w_{11}^{-1} w_{12}, \quad w_{22} = w_{21} X = w_{21} w_{11}^{-1} w_{12}.

$

Please refer to [4] for the detailed proof.

Weaknesses 3: The main text requires rigorous definitions of notations, as it is difficult to follow the proofs of the main results without them.

Answer: The following table presents the definitions of all the notations used in the submission.

Table 2: Definitions of Notations

Table 1: Client and Server Ranks Over 25 Rounds on LLaMA 7B+MPQA

It is obvious that $\Delta \tilde{\mathbf{W}} \neq \Delta \mathbf{W}$ when $\mathbf{U}^{1} \neq \mathbf{U}^{2}$, $\mathbf{\Sigma}^{1} \neq \mathbf{\Sigma}^{2}$, and $\mathbf{V}^{1} \neq \mathbf{V}^{2}$ due to the data heterogeneity property of federated learning. Therefore, the client-drift issue is raised.

The difference between $\Delta \tilde{\mathbf{W}}$ and $\Delta \mathbf{W}$ is mainly due to the noise introduced by the cross-products of LoRA modules from different clients. This difference may become more significant when (1) the number of local update steps between aggregations is large and (2) the local datasets are different across clients.

We propose a Riemannian parameter matching method to match the local parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ on other clients with pivots $\mathbf{U}^1$ and $\mathbf{V}^1$ on client 1, in terms of their lengths and directions, i.e. $\tilde{\mathbf{U}}^k = \mathbf{U}^k \mathbf{R}^k \approx \mathbf{U}^1$ and $\tilde{\mathbf{V}}^k = \mathbf{S}^k \mathbf{V}^k \approx \mathbf{V}^1$. To maintain the consistency between low-rank parameter metrics before and after the Riemannian parameter matching, i.e., $\mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2 = \tilde{\mathbf{U}}^2 \tilde{\mathbf{\Sigma}}^2 \tilde{\mathbf{V}}^2 \approx \mathbf{U}^1 \tilde{\mathbf{\Sigma}}^2 \mathbf{V}^1$, for ensuring the effectiveness of FFT-FM with rank-adaptive LoRA, we derive a modified diagonal matrix $\tilde{\mathbf{\Sigma}}^k$ for the other clients by performing the SVD on low-dimensional $r \times r$ matrices $(\mathbf{U}^1)^T \mathbf{U}^k$ and $(\mathbf{V}^1)^T \mathbf{V}^k$. Thus, the client-drift issue is resolved.

$

\begin{split} \Delta \tilde{\mathbf{W}} &=\frac{1}{2}\left(\mathbf{U}^{1}+ \tilde{\mathbf{U}^2}\right) \times \frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right) \times \frac{1}{2}\left(\mathbf{V}^{1}+ \tilde{\mathbf{V}^{K}}\right) \ &\approx\mathbf{U}^{1} \times \frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right) \times \mathbf{V}^{1} = \frac{1}{2}\left(\mathbf{U}^1 \mathbf{\Sigma}^1 \mathbf{V}^1 + \mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2\right) = \Delta \mathbf{W} \end{split}

$

Based on the global diagonal matrix $\frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right)$, it is easy to find the optimal rank of the global parameter matrix, with aggregation on only the local low-rank matrices $\mathbf{U}^k, \tilde{\mathbf{\Sigma}}^{k}$, and $\mathbf{V}^k$.

Rank-drift issue: The data heterogeneity is a common challenging problem in the federated learning. The collected data on different clients are Non-Independent Identically Distributed (Non-IID). Thus, the parameters of local models trained on the non-IID data are quite different from each other as well as the parameters of global model. Each local model may oscillate back and forth in different training rounds, leading to unstable and slow convergence and causing suboptimal model performance.

At the same time, due to the heterogeneity of local models, the optimal ranks of local models are quite different from each other as well as the one of global model, raising the rank-drift issue. Similarly, the oscillation of the optimal rank of each local model can slow down the model convergence and degrade the model performance. For example, given two clients 1 and 2, in round 1, the local model on client 1/2 has the optimal rank 5/20. The server aggregates local models from clients to generate a global model based on FedAvg. The global model by model aggregation has the optimal rank 12. In round 2, after the global model are sent back to the clients, the clients begin the training with the global model with the optimal rank 12. Thus, the optimal rank of local models oscillate between 5/20 and 12. The rank oscillation may repeat in each training round, slowing down the model convergence.

The following table presents a case study regarding the rank-drift in federated learning. In round 1, the optimal rank of the global model is 16. In round 2, the optimal rank is changed to 11. In subsequent rounds, the optimal rank continue to oscillate back and forth. The ranks of local models fluctuate around values such as 9.75, 10.17, and 10.53, while the rank of global modelgradually decreases and stabilizes at 8. These results confirm that rank-drift introduces instability into the training process, slowing convergence and degrading performance, which emphasizes the importance of our Riemannian gradient descent optimization approach to stabilize ranks and ensure efficient training.

We thank this reviewer for the commnets.

Weaknesses 1: The writing quality is poor, with too many long sentences that are difcult to follow. Additionally, the paper includes many unexplained abbreviations, such as FedAvg, RAFFT, and RGD, which appear in the introduction without clarification.

Answer: Thank you for your valuable feedback on our work. The following table 1 presents the definitions of all the abbreviations used in the submission.

Table 1: List of abbreviations used in the paper

Weaknesses 2: The authors need to provide clear definitions of the client-drift and rank-drift issues.

Answer: Client-drift issue: The client-drift issue is due to the aggregation of low-rank matrices in the FFT-FM. Here, we introduce an illustrative example with two clients to better understand the client-drift issue.

Concretely, if the clients locally fine-tune on full parameters and the server aggregates with FedAvg, the new global model parameters can be expressed as follows.

$

\mathbf{W} = \mathbf{W}^0 +\frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2)

$

In our work, the clients locally fine-tune on both low-rank parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ and diagonal matrix $\mathbf{\Sigma}^k$ to determine the rank. The server uses FedAvg to aggregate the low-rank matrices and diagonal matrix. The local model updates $\Delta \mathbf{W}^k$ are represented as:

$

\Delta \mathbf{W}^k = \mathbf{U}^k \mathbf{\Sigma}^k \mathbf{V}^k,  \quad k ∈ \\{1, 2\\}

$

The expected global model parameter updates $\Delta \mathbf{W}$ can be expressed as follows.

$

\Delta \mathbf{W} = \frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2) = \frac{1}{2} (\mathbf{U}^1 \mathbf{\Sigma}^1 \mathbf{V}^1 + \mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2)

$

After using FedAvg to aggregate the local matrices $\mathbf{U}^k, \mathbf{\Sigma}^k, \mathbf{V}^k$, the server produces

where the underlined terms denote the difference between $\Delta \tilde{\mathbf{W}}$ and $\Delta \mathbf{W}$.

Thus, we propose a Riemannian parameter matching method to match the local parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ on other clients with pivots $\mathbf{U}^1$ and $\mathbf{V}^1$ on client 1, in terms of their lengths and directions, i.e. $\tilde{\mathbf{U}}^k = \mathbf{U}^k \mathbf{R}^k \approx \mathbf{U}^1$ and $\tilde{\mathbf{V}}^k = \mathbf{S}^k \mathbf{V}^k \approx \mathbf{V}^1$. To maintain the consistency between low-rank parameter metrics before and after the Riemannian parameter matching, i.e., $\mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2 = \tilde{\mathbf{U}}^2 \tilde{\mathbf{\Sigma}}^2 \tilde{\mathbf{V}}^2 \approx \mathbf{U}^1 \tilde{\mathbf{\Sigma}}^2 \mathbf{V}^1$, for ensuring the effectiveness of FFT-FM with rank-adaptive LoRA, we derive a modified diagonal matrix $\tilde{\mathbf{\Sigma}}^k$ for the other clients by performing the SVD on low-dimensional $r \times r$ matrices $(\mathbf{U}^1)^T \mathbf{U}^k$ and $(\mathbf{V}^1)^T \mathbf{V}^k$. Thus, the client-drift issue is resolved.

$

\begin{split} \Delta \tilde{\mathbf{W}} &=\frac{1}{2}\left(\mathbf{U}^{1}+ \tilde{\mathbf{U}^2}\right) \times \frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right) \times \frac{1}{2}\left(\mathbf{V}^{1}+ \tilde{\mathbf{V}^{K}}\right) \ &\approx\mathbf{U}^{1} \times \frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right) \times \mathbf{V}^{1} = \frac{1}{2}\left(\mathbf{U}^1 \mathbf{\Sigma}^1 \mathbf{V}^1 + \mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2\right) = \Delta \mathbf{W} \end{split}

$

Based on the global diagonal matrix $\frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right)$, it is easy to find the optimal rank of the global parameter matrix, with aggregation on only the local low-rank matrices $\mathbf{U}^k, \tilde{\mathbf{\Sigma}}^{k}$, and $\mathbf{V}^k$.

Weakness 2: Also, Reiemannian LoRA has been studied in a centralized setting in the following paper: Riemannian Preconditioned LoRA for Fine-Tuning Foundation Models, ICML 2024.

Answer: The research objective of this ICML 2024 paper is completely different from ours. It aims to use different learning rates on the low-rank parameter matrices $A$ and $B$, introducing an $r \times r$ preconditioner to stabilize the LoRA learning by scaling gradients. This gradient scaling method can be derived from a new Riemannian metric. This paper does not consider the rank-drift issue and address the optimization on a Riemannian manifold and maintaining a fixed rank during updates. However, one of main contributions in this paper is to address the rank-drift issue, which is unique in the FFT-FM but not in centralized learning. Therefore, the centralized learning method in the ICML 2024 paper cannot be directly applied to our scenario.

Rank-drift issue: The data heterogeneity is a common challenging problem in the federated learning. The collected data on different clients are Non-Independent Identically Distributed (Non-IID). Thus, the parameters of local models trained on the non-IID data are quite different from each other as well as the parameters of global model. Each local model may oscillate back and forth in different training rounds, leading to unstable and slow convergence and causing suboptimal model performance.

At the same time, due to the heterogeneity of local models, the optimal ranks of local models are quite different from each other as well as the one of global model, raising the rank-drift issue. Similarly, the oscillation of the optimal rank of each local model can slow down the model convergence and degrade the model performance. For example, given two clients 1 and 2, in round 1, the local model on client 1/2 has the optimal rank 5/20. The server aggregates local models from clients to generate a global model based on FedAvg. The global model by model aggregation has the optimal rank 12. In round 2, after the global model are sent back to the clients, the clients begin the training with the global model with the optimal rank 12. Thus, the optimal rank of local models oscillate between 5/20 and 12. The rank oscillation may repeat in each training round, slowing down the model convergence.

We thank this reviewer for the helpful comments.

Weakness 1: Many of the presentations of the paper are similar with the ones in the (Sun et al., 2024) paper. As an example, equation 1 in the authors's paper is also presented in equation 3 of (Sun et al., 2024). It is not clear which part is new to this work. I feel that the authors need to add more details regarding the difference with the previous work.

Answer: The problem, challenges, and techniques in our paper significantly differ from the above ICLR 2024 (Sun et al., 2024), although both papers have the client-drift issue and the reason of the issue in both papers are the same, due to the aggregation of low-rank matrices in the FFT-FM.

The ICLR 2024 paper conducts this issue raised by only low-rank parameter matrices $\mathbf{A}^k$ and $\mathbf{B}^k$, but our work addresses a more challenging problem by considering both low-rank parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ and diagonal matrix $\mathbf{\Sigma}^k$ to determine the rank. Here, we introduce an illustrative example with two clients to better understand the difference of the client-drift issues and the corresponding techniques between the ICLR 2024 paper and our paper.

Concretely, if the clients locally fine-tune on full parameters and the server aggregates with FedAvg, the new global model parameters can be expressed as follows.

$

\mathbf{W} = \mathbf{W}^0 +\frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2)

$

In the ICLR 2024 paper, when the clients locally fine-tune on low-rank parameters based on LoRA and the server uses FedAvg to aggregate the low-rank matrices, the global model parameters can be expressed as follows.

$

\underbrace{\tilde{\mathbf{W}} = \mathbf{W}^0 + \frac{1}{2}(\mathbf{B}^1 + \mathbf{B}^2) \times \frac{1}{2}(\mathbf{A}^1 + \mathbf{A}^2)}_{\text{Parameter aggregation with LoRA + FedAvg}} \neq \underbrace{\mathbf{W}^0 + \frac{1}{2}(\mathbf{B}^1 \mathbf{A}^1 + \mathbf{B}^2 \mathbf{A}^2) = \mathbf{W}_0 +\frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2) = \mathbf{W}}_{\text{Ideal parameter aggregation with FedAvg on full parameter matrices}}

$

When the clients use LoRA locally, we have $\Delta \mathbf{W}^k = \mathbf{B}^k \mathbf{A}^k$ on client $k$. The data heterogeneity is a common challenging problem in the federated learning. The collected data on different clients are Non-Independent Identically Distributed (NonIID). Thus, the parameters of local models trained on the non-IID data are quite different from each other, i.e., $\mathbf{B}^1 \neq \mathbf{B}^2$ and $\mathbf{A}^1 \neq \mathbf{A}^2$. This leads to the client-drift issue in the ICLR 2024 paper, i.e., the difference between $\mathbf{W}$ and $\tilde{\mathbf{W}}$ is not equal to 0. In order to address the client-drift issue, the ICLR 2024 paper keeps both $\mathbf{W}^0$ and $\mathbf{A}^0$ frozen and makes only $\mathbf{B}$ trainable.

$

\underbrace{\tilde{\mathbf{W}} = \mathbf{W}^0 + \frac{1}{2}(\mathbf{B}^1 + \mathbf{B}^2) \times \mathbf{A}^0}_{\text {Parameter aggregation with LoRA + FedAvg}} = \underbrace{\mathbf{W}^0 + \frac{1}{2}(\mathbf{B}^1 \mathbf{A}^0 + \mathbf{B}^2 \mathbf{A}^0) = \mathbf{W}_0 +\frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2) = \mathbf{W}}_{\text {Ideal parameter aggregation with FedAvg on full parameter matrices}}

$

In our work, the clients locally fine-tune on both low-rank parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ and diagonal matrix $\mathbf{\Sigma}^k$ to determine the rank. The server uses FedAvg to aggregate the low-rank matrices and diagonal matrix. The local model $\Delta \mathbf{W}^k$updates are represented as:

$

\Delta \mathbf{W}^k = \mathbf{U}^k \mathbf{\Sigma}^k \mathbf{V}^k, \quad k \in \\{1, 2\\}

$

The expected global model parameter updates $\Delta \mathbf{W}$ can be expressed as follows.

$

\Delta \mathbf{W} = \frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2) = \frac{1}{2} (\mathbf{U}^1 \mathbf{\Sigma}^1 \mathbf{V}^1 + \mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2)

$

Question 8: How are the projected gradients in section 4 computed? Could the authors clarify the advantages compared to Zhang et al. 2023, conventional LoRA, and full model fine-tuning? This may be related to the definition of rank-drift issue, but I am not sure.

Answer: Before explaining projected gradients, we first compare RAFFT with previous works. Methods such as AdaLoRA (Zhang et al., 2023), conventional LoRA, and full model fine-tuning are designed for single-machine settings. When extended to federated learning scenarios, these methods encounter the rank-drift issue due to data heterogeneity, as described in detail in the response to Question 4. However, these methods do not guarantee that the selected optimal rank remains consistent during updates, negatively impacting convergence speed and model performance. In contrast, our approach utilizes a Riemannian manifold encompassing all fixed-rank matrices. By optimizing directly on this manifold, we ensure that the rank remains invariant throughout client updates. This is because the Riemannian manifold inherently constrains the optimization process to matrices of the same fixed rank. In contrast, traditional gradient descent methods like SGD have no such constraints and may cause updates to deviate from the manifold, resulting in unintended changes in rank.

In contrast to traditional optimization techniques, our approach ensures rank invariance during updates by leveraging the geometry of the Riemannian manifold. To achieve this, we project the Euclidean gradient onto the tangent space $T_{M}$ at the point $\Delta \mathbf{W}$ on the manifold $M$, thereby obtaining the Riemannian gradient. The projected gradient lies within the tangent space, ensuring that the update direction respects the manifold's constraints.

After the gradient step, the updated point is retracted back onto the Riemannian manifold using a retraction operation, which guarantees that the updated model parameters remain within the space of fixed-rank matrices. By keeping updates confined to the Riemannian manifold, our approach preserves the rank consistency throughout local updates, addressing the limitations of conventional methods that may inadvertently alter the rank during optimization.

Therefore, our method alleviates the rank-drift issue to speed up the convergence of RAFFT.

For specific formulas and proofs, please refer to the appendix. The projection function is provided in Appendix 7.3 (Eq. 61). The Riemannian gradient is defined in Section 4 (Eq. 21), and the retraction function is described in Appendix 7.3 (Eq. 68). Furthermore, the process is detailed in Algorithm 1 (lines 11 to 14), Appendix 7.9.

Question 9: Minor issue: Some references are duplicated, e.g., Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficient fine-tuning. ICLR’23 Youbang Sun, Zitao Li, Yaliang Li, and Bolin Ding. Improving lora in privacy-preserving federated Learning. ICLR’24

Answer: Thanks for your helpful comments. We have carefully proofread our manuscript and corrected the duplication of the mentioned references in the rebuttal revision to enhance the paper's readability.

Question 5: In section 3, the content seems to be a mix of derivation and algorithms. It is unclear what exactly needs to be done on both clients and the server.

Answer: In our paper's appendix (Section 7.9, Algorithm 1), we have provided a detailed description of the complete RAFFT algorithm. In each communication round, clients perform local updates using the Riemannian gradient and the retraction function, ensuring that the optimal rank selected by the server remains unchanged (lines 11-14). After the updates, each client sends the updated parameters back to the server (line 17). The server then performs Riemannian parameter matching between the $\mathbf{U}$ and $\mathbf{V}$ matrices of all clients and those of the first client to achieve alignment in both length and direction (line 24-28). The server adjusts the corresponding singular value matrix to ensure consistency before and after matching (line 29). The aggregated $\mathbf{U}$, $\mathbf{V}$, and $\mathbf{\Sigma}$ matrices from all clients are then combined (line 30-33). The rank for the next round of training is determined based on the values in the singular value matrix, and the model parameter matrix is updated accordingly (line 34-37). The server distributes the updated $\mathbf{U}$, $\mathbf{V}$, and $\mathbf{\Sigma}$ matrices back to the clients, and training continues until convergence.

Question 6: It is unclear to me what the advantage of Riemann aggregation in section 3 is compared to aggregating the reconstructed \Delta W from every client.

Answer: Thanks for your helpful comments. Reconstructing $\Delta \mathbf{W}$ requires every client to perform full-matrix SVD decomposition to select the optimal rank, which is computationally expensive. In contrast, our Riemann aggregation method eliminates this overhead by allowing the server to directly determine the rank during aggregation. Instead of reconstructing $\Delta \mathbf{W}$, our approach operates directly in the decomposed matrix space ($\mathbf{U}$, $\mathbf{\Sigma}$, and $\mathbf{V}$), where the singular values in $\mathbf{\Sigma}$ are leveraged to dynamically adapt the rank.

Question 7: It is unclear how LoRA (SVD) parameters are determined. It is unclear whether adaptive ranks (what kind of adaptive?), or fixed rank are used.

Answer: We initialize the parameters following the previous work, AdaLoRA. Specifically, we randomly initialize the matrices $\mathbf{U}$ and $\mathbf{V}$, and set $\mathbf{\Sigma}$ to 0, ensuring that $\Delta \mathbf{W}=0$ at the beginning of training. During subsequent parameter updates, all three matrices are updated.

Our overall goal is to achieve adaptive rank LoRA fine-tuning in a federated learning scenario by combining dynamic rank selection on the server side with fixed-rank optimization on the client side. This two-part strategy ensures that the federated system remains stable and effectively adapts to the heterogeneity of client data distributions, improving both performance and convergence.

Specifically:

On the server side, adaptive rank selection dynamically determines the optimal rank after each training round. Our approach uses the singular value contribution rate, $\Theta(r)$, to select the rank $\tilde{\mathbf{\Sigma}}^{k'}$. The contribution rate $\Theta(r)$ is defined as the cumulative percentage of the $r$ largest singular values relative to the total sum of all singular values. The server retains the $r$ singular values such that $\Theta(r)$ exceeds a predefined threshold $\varphi$. For instance, if there are 16 singular values in total and the cumulative sum of the 5 largest singular values accounts for more than 90% of the total, the server adjusts the rank to 5. This method ensures that the server retains the most important information while reducing redundancy in the parameter matrices.

On the client side, by leveraging Riemannian manifold theory, we develop a Riemannian gradient descent (RGD) method to guarantee the local full parameter matrices on clients in the form of low-rank ones with fixed rank optimized by the server in each FFT-FM round.

After using FedAvg to aggregate the local matrices $\mathbf{U}^k, \mathbf{\Sigma}^k, \mathbf{V}^k$, the server produces

where the underlined terms denote the difference between $\Delta \tilde{\mathbf{W}}$ and $\Delta \mathbf{W}$.

It is obvious that $\Delta \tilde{\mathbf{W}} \neq \Delta \mathbf{W}$ when $\mathbf{U}^{1} \neq \mathbf{U}^{2}$, $\mathbf{\Sigma}^{1} \neq \mathbf{\Sigma}^{2}$, and $\mathbf{V}^{1} \neq \mathbf{V}^{2}$ due to the data heterogeneity property of federated learning. Therefore, the client-drift issue is raised.

The difference between $\Delta \tilde{\mathbf{W}}$ and $\Delta \mathbf{W}$ is mainly due to the noise introduced by the cross-products of LoRA modules from different clients. This difference may become more significant when (1) the number of local update steps between aggregations is large and (2) the local datasets are different across clients.

We propose a Riemannian parameter matching method to match the local parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ on other clients with pivots $\mathbf{U}^1$ and $\mathbf{V}^1$ on client 1, in terms of their lengths and directions, i.e. $\tilde{\mathbf{U}}^k = \mathbf{U}^k \mathbf{R}^k \approx \mathbf{U}^1$ and $\tilde{\mathbf{V}}^k = \mathbf{S}^k \mathbf{V}^k \approx \mathbf{V}^1$. To maintain the consistency between low-rank parameter metrics before and after the Riemannian parameter matching, i.e., $\mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2 = \tilde{\mathbf{U}}^2 \tilde{\mathbf{\Sigma}}^2 \tilde{\mathbf{V}}^2 \approx \mathbf{U}^1 \tilde{\mathbf{\Sigma}}^2 \mathbf{V}^1$, for ensuring the effectiveness of FFT-FM with rank-adaptive LoRA, we derive a modified diagonal matrix $\tilde{\mathbf{\Sigma}}^k$ for the other clients by performing the SVD on low-dimensional $r \times r$ matrices $(\mathbf{U}^1)^T \mathbf{U}^k$ and $(\mathbf{V}^1)^T \mathbf{V}^k$. Thus, the client-drift issue is resolved.

$

\begin{split} \Delta \tilde{\mathbf{W}} &=\frac{1}{2}\left(\mathbf{U}^{1}+ \tilde{\mathbf{U}^2}\right) \times \frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right) \times \frac{1}{2}\left(\mathbf{V}^{1}+ \tilde{\mathbf{V}^{K}}\right) \ &\approx\mathbf{U}^{1} \times \frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right) \times \mathbf{V}^{1} = \frac{1}{2}\left(\mathbf{U}^1 \mathbf{\Sigma}^1 \mathbf{V}^1 + \mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2\right) = \Delta \mathbf{W} \end{split}

$

Based on the global diagonal matrix $\frac{1}{2}\left(\mathbf{\Sigma}^{1}+ \tilde{\mathbf{\Sigma}}^{2}\right)$, it is easy to find the optimal rank of the global parameter matrix, with aggregation on only the local low-rank matrices $\mathbf{U}^k, \tilde{\mathbf{\Sigma}}^{k}$, and $\mathbf{V}^k$.

We thank this reviewer for the constructive comments.

Weakness 1: I also find it a bit hard to interpret the results on why the proposed RAFFT methods are better, and the difference of the metrics (accuracy, loss and time)

Answer: Our paper addresses two unique challenges in the FFT-FM: client drift and rank drift.

First, by leverages Riemannian parameter matching to dynamically adjust the rank based on the singular value matrix $\mathbf{\Sigma}$ after aggregating global parameters on the server side. This approach eliminates the need for additional operations such as full matrix SVD decomposition. Furthermore, it resolves the client-drift issue, which arises from independently aggregating the client parameter matrices $\mathbf{U}^k$, $\mathbf{\Sigma}^k$, and $\mathbf{V}^k$.

Second, by leveraging Riemannian manifold theory, we develop a Riemannian gradient descent (RGD) method to guarantee the local full parameter matrices on clients in the form of low-rank ones with fixed rank optimized by the server in each FFT-FM round, for alleviating the rank-drift issue to speed up the convergence of RAFFT.

In line with established protocols from prior studies on federated learning for large language models [1, 2, 3], we evaluate the efficacy of parameter-efficient fine-tuning for classification tasks using three key metrics: Loss, Accuracy, and Time. Loss quantifies the model's prediction error, providing insights into the effectiveness of the learning process. Accuracy assesses the model's performance in classifying unseen data, reflecting the practical applicability of the federated learning model. Time measures the speed of convergence and the overall computational demand, both of which are critical in federated settings where computational resources and time are often constrained.

In our submission, we have conducted a comprehensive analysis of our proposed method, RAFFT, across six datasets and against 13 baseline methods. The experimental results, detailed in Tables 1-2 and 4-13 of the submission, demonstrate the significant advantages of RAFFT. Specifically, RAFFT achieves notable accuracy improvements, with a maximum increase of 30.22% over other baseline methods. Moreover, RAFFT excels in extreme non-IID data environments, achieving an accuracy of up to 58.05%, highlighting its robustness under challenging data distributions. In terms of efficiency, RAFFT reduces training time by up to 41.86%, underscoring its practicality and effectiveness in federated learning scenarios where both accuracy and computational efficiency are critical.

REFERENCES:

[1] Sara Babakniya, Ahmed Elkordy, Yahya Ezzeldin, Qingfeng Liu, Kee-Bong Song, MOSTAFA EL- Khamy, and Salman Avestimehr. SLoRA: Federated parameter efficient fine-tuning of language models. In International Workshop on Federated Learning in the Age of Foundation Models in Conjunction with NeurIPS 2023, 2023.

[2] Yae Jee Cho, Luyang Liu, Zheng Xu, Aldi Fahrezi, Matt Barnes, and Gauri Joshi. Heterogeneous loRA for federated fine-tuning of on-device foundation models. In International Workshop on Federated Learning in the Age of Foundation Models in Conjunction with NeurIPS 2023, 2023.

[3] Youbang Sun, Zitao Li, Yaliang Li, and Bolin Ding. Improving loRA in privacy-preserving federated learning. In The Twelfth International Conference on Learning Representations, 2024b.

Question 1: I kindly ask the authors to clarify "the client-drift issue" for LoRA. Feel free to use the definition in Sun et al. 2024a and Wang et al. 2024, but I would make my decision based on a self-contained explanation. Eq (1) is a fact, but it is not clear to me if it is an issue.

Answer: Client-drift issue: The client-drift issue is due to the aggregation of low-rank matrices in the FFT-FM. Here, we introduce an illustrative example with two clients to better understand the client-drift issue.

Concretely, if the clients locally fine-tune on full parameters and the server aggregates with FedAvg, the new global model parameters can be expressed as follows.

$

\mathbf{W} = \mathbf{W}^0 +\frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2)

$

In our work, the clients locally fine-tune on both low-rank parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$ and diagonal matrix $\mathbf{\Sigma}^k$ to determine the rank. The server uses FedAvg to aggregate the low-rank matrices and diagonal matrix. The local model updates $\Delta \mathbf{W}^k$ are represented as:

$

\Delta \mathbf{W}^k = \mathbf{U}^k \mathbf{\Sigma}^k \mathbf{V}^k,  \quad k ∈ \\{1, 2\\}

$

The expected global model parameter updates $\Delta \mathbf{W}$ can be expressed as follows.

$

\Delta \mathbf{W} = \frac{1}{2} (\Delta \mathbf{W}^1+\Delta \mathbf{W}^2) = \frac{1}{2} (\mathbf{U}^1 \mathbf{\Sigma}^1 \mathbf{V}^1 + \mathbf{U}^2 \mathbf{\Sigma}^2 \mathbf{V}^2)

$

Question 11: Even if we acknowledge these issues and try to resolve them, I am not fully convinced the advantage of the proposed method. For example, why the proposed method is better than the following simple strategy: run LoRA on each client, reconstruct $\Delta W$ and send back to server, server aggregate $\Delta W$ from clients, and run SVD before distributing LoRA module with different rank to each client. The $\Delta W$ reconstruction can happen on the server, though that potentially raises more privacy concerns.

Answer: As mentioned in early discussion, Cho et al. (2024) [7] observed that reconstruction $\Delta W$ before aggregation achieves worse performance in their experiments. Another Mahla et al. [8] work also reported that reconstructing $\Delta \mathbf{W}$ before aggregation leads to performance degradation.

This kind of method based on the reconstruction of $\Delta \mathbf{W}$ before aggregation will result in the following three drawbacks, in terms of effectiveness and efficiency:

1. Performance Degradation due to Matrix Reconstruction:

Reconstructing matrices before aggregation leads to performance degradation, as it causes the loss of critical information and features during the aggregation process. This results in incomplete updates and negatively impacts model performance, as observed in Cho et al. (2024) and Mahla et al. , where the reconstruction process contributed to a significant drop in accuracy (See their results in Tables 3 in [7] and Tables 1-2 in [8], [7] has an average of 46% accuracy degrade and [8] has an average of 73% F1 score decrease).

2. Excessive Memory and Communication Overhead:

Reconstructing $\Delta \mathbf{W}$ creates substantial communication and memory costs, making it impractical for large-scale federated learning. The LLaMA-7B model is used in our experiments, reconstructing $\Delta \mathbf{W}$ will produce a high-dimensional full parameter matrix with 6,742,618,112 elements in total. In the setting of federated learning, each client needs to upload this high-dimensional full parameter matrix to the server, resulting in a high communication cost of approximately 6.75GB per client when using float32 precision (6,742,618,112 * 4bytes). In addition, the server needs huge memory consumption to maintain these local reconstructed $\Delta \mathbf{W}$. This expensive strategy is impractical and not scalable for high-dimensional foundation models. Current mainstream methods of federated learning of foundation models is to directly aggregate on low-rank parameter matrices without the reconstruction [1-8]. Our work follows the same approaches.

3. High Computational Complexity of SVD:

After reconstructing $\Delta \mathbf{W}$, performing singular value decomposition (SVD) on the global full parameter matrix introduces substantial computational overhead. For a model like LLaMA-7B, where the parameter matrix dimensions are $4096 \times 4096$, the computational complexity of SVD is $O(n^3)$, which means approximately $4096^3$ ($\sim 68.72$ billion floating-point operations) per round. This high computational burden makes the method inefficient and impractical for large-scale federated learning of foundation models.

In contrast, our method avoids the reconstruction of $\Delta \mathbf{W}$ by design with improved effectiveness and efficiency:

1. Avoiding Matrix Reconstruction: We do not perform the reconstruction of full parameter matrices. Consequently, our method does not suffer from the information loss associated with these approaches. By preserving the complete low-rank structure of the parameter matrices, our method ensures both efficiency and effectiveness in federated learning.

2. Reduced Communication and Memory Overhead:
   Clients in our method directly upload independent low-rank matrices, significantly reducing communication and memory costs. For the LLaMA-7B model, our method only requires a communication volume of $78,962 \times 4 \, \text{bytes}$ $(\sim 0.3 \, \text{MB})$ per client. This represents a significant reduction compared to the 6.75GB communication cost caused by reconstructing $\Delta W$.

3. Dynamic Rank Selection without SVD on Full-parameter Matrix:
   In our method, the server selects the optimal rank based on the singular value matrix $\mathbf{\Sigma}$ without performing SVD on the full parameter matrix, resulting in reduced computational complexity. Specifically,we perform SVD on low-dimensional $r \times r$ matrices, where $r$ is the rank parameter in LoRA. In the case of $r = 8$ in our experiments, our computational complexity $8^3$ is significantly reduced compared to that of the full matrix SVD $4096^3$.

In addition, our experimental results have validated the effectiveness and efficiency of our method for federated fine-tuning of foundation models (FFT-FM).

Question 10: while I appreciate the experiments and observation that ranks are different on server and clients, I believe both client-drift and rank-drift need clearer definition. The current description are "facts" rather than "issues", and I am not fully convinced.

Answer: Thank you for your follow-up feedback. The client-drift issue is defined as the discrepancy between independently aggregating the local low-rank parameter matrices on clients and aggregating the full parameter matrices on clients. Specifically, when the server aggregates low-rank parameter matrices $\frac{1}{k}\Sigma_k\mathbf{U}^k \times \frac{1}{k}\Sigma_k\mathbf{\Sigma}^k \times \frac{1}{k}\Sigma_k\mathbf{V}^k$ independently, the resulting updates deviate from the exact centralized LoRA $\mathbf{U} \times \mathbf{\Sigma} \times \mathbf{V}$ as well as the aggregation $\frac{1}{k}\Sigma_k\mathbf{U}^k \times \mathbf{\Sigma}^k \times \mathbf{V}^k$ of the full parameter matrices on clients. We have provide the detailed analysis and discussion the discrepancy between two aggregation methods in the setting of federated learning in the answer to Question 1. The main reason is due to the noise introduced by the cross-products of the local low-rank parameter matrices from different clients. This discrepancy may become more significant when (1) the number of local update steps between aggregations is large and (2) the local datasets are different across clients.

The client-drift issue has been clearly defined in previous studies including [Sun et al. (2024a) [1]; Wang et al. (2024) [2]; Singha et al. (2024) [3]; Guo et al. (2024) [4]; Koo et al. (2024) [5]; Nguyen et al. (2024)] [6], where they explicitly defined the client-drift problem caused by the aggregation of the local low-rank parameter matrices in federated fine-tuning of foundation models (FFT-FM), similar to the definition in our paper. All these previous works suggest that the client drift issue is not just a fact but a critical issue that needs to be addressed. Our definition of the client-drift problem is consistent with these works, as the reason for the issue in all cases stems from the aggregation of low-rank matrices in FFT-FM.

In addition, our work addresses a more complex problem. While the aforementioned papers limit their focus to low-rank parameter matrices $\mathbf{A}_k$ and $\mathbf{B}_k$, our method simultaneously considers the aggregation of three components: low-rank parameter matrices $\mathbf{U}^k$ and $\mathbf{V}^k$, as well as the diagonal matrix $\mathbf{\Sigma}^k$, which determines the rank. By addressing this more challenging problem, our approach goes beyond the scenarios discussed in prior works, providing a more comprehensive solution to the client-drift issue in FFT-FM.

Rank-drift issue refers to the discrepancy between the global optimal rank on the server and the locally optimal ranks on individual clients, caused by data heterogeneity, leading to oscillations in the optimal ranks during training and slowing down the convergence.

Data heterogeneity is a common and challenging problem in federated learning systems. Let $r^\star$ denote the global optimal rank on the server and $r^k$ represent the optimal rank for each client. Due to the data heterogeneity, these local optimal ranks $r^k$ may differ significantly from each other and from the global optimal rank $r^\star$. Each client, influenced by its non-IID data, pushes the model toward its own optimal rank in the optimization space, moving away from the global optimal rank. Even if all clients start with the same rank, they will eventually converge toward their individual optimal ranks $r^k$, causing $r^k$ to deviate from $r^\star$.

This discrepancy results in rank-drift, where the optimal rank for each client oscillates back and forth across training rounds. Consequently, this rank deviation also causes the local parameter matrices to deviate from the optimal global parameter matrices, leading to unstable and slow convergence, ultimately resulting in suboptimal model performance. The issue is particularly pronounced in scenarios where client models are trained on highly heterogeneous data, as the varying ranks exacerbate the divergence from the global optimum.

For example, given two clients 1 and 2, in round 1, the local model on client 1/2 has the optimal rank 5/20. The server aggregates local models from clients to generate a global model based on FedAvg. The global model by model aggregation has the optimal rank 12. In round 2, after the global model are sent back to the clients, the clients begin the training with the global model with the optimal rank 12. Thus, the optimal rank of local models oscillate between 5/20 and 12. The rank oscillation may repeat in each training round, slowing down the model convergence.