Towards Robust and Efficient Federated
Low-Rank Adaptation with Heterogeneous Clients

LoRA approximates the fine-tuned weight as \( W = W_0 + \Delta W = W_0 + BA \), where \( W_0 \) is the frozen pre-trained matrix, and \( B \in \mathbb{R}^{d_1 \times r}, A \in \mathbb{R}^{r \times d_2} \) are low-rank matrices with rank \( r \ll \min\{d_1, d_2\} \). This reduces the number of trainable parameters from \( d_1 \cdot d_2 \) to \( r \cdot (d_1 + d_2) \), which benefits in both communication and computation efficiency in Federated Learning (FL).

However, due to the bilinear parametrization of LoRA, aggregating client updates \( ( \Delta W_k = B_kA_k \; \forall k \in [K] )\) leads to the discordance problem:

where \( w_k \)'s are non-negative aggregation weights with \( \sum w_k = 1 \). This mismatch degrades performance, which has led to a growing body of research aimed at addressing this issue.

• FFA-LoRA [Sun et al.] Let \( A_k = A_0 \) be frozen for all \( k \in [K] \), so that
  \[ \sum_{k=1}^K w_k B_k A_k = \left( \sum_{k=1}^K w_k B_k \right) A_0. \]
• FlexLoRA [Bai et al.] Aggregate \( B_kA_k \) directly. Then, using SVD, decompose it back into
  \[ \sum_{k=1}^K w_k B_k A_k = BA. \]

Instead of solely training module \( B \) while keeping module \( A \) frozen permanently, LoRA-A² alternates between the two: LoRA module \( B \) is frozen during even rounds, while module \( A \) is frozen during odd rounds. This method preserves the optimization space while effectively resolving discordance. Specifically, when freezing \( A \), we have

and when freezing \( B \), we have In this way, LoRA-A² trains both \( B \) and \( A \), ensuring that \( A \) does not remain the same as its initial value. (Note: Since the standard convention of LoRA initializes \( B \) to zero, we first train \( B \). However, if a different initialization scheme is used, training \( A \) first could also be a reasonable choice.)

LoRA-A² selects important LoRA ranks to match local communication rank budget \( r_i \) out of global LoRA adapter with rank \( r_G \) adaptively based on the local dataset. This approach provides two key benefits:

• It minimizes client conflicts by allowing different clients to choose different LoRA ranks in high heterogeneity.
• It reallocates rank resources from unimportant LoRA modules to modules that require more fine-tuning, which is especially effective when communication rank budget is small.

Figure 2 below shows the number of selected ranks per module for three clients in a pathological data distribution setting. In this setup, clients 0 and 1 are assigned the "medical" and "space" classes respectively, while client 2 is assigned "motorcycle" and "religions" classes from the 20 Newsgroups dataset. This indicates that clients with semantically related data tend to converge on similar rank subspaces, potentially enabling more effective cooperative training. Meanwhile, clients with divergent data distributions select distinct ranks, leading to more independent updates. In addition, we observe that most modules are assigned zero ranks across clients. This suggests that our adaptive mechanism effectively prunes out modules that do not benefit from fine-tuning. Please refer to Sections 4.2 and 4.3 for the detailed methodology of our rank selection mechanism.

Figure 2: Visualization on number of selected rank per module. The x-axis shows RoBERTa module types, while the y-axis indicates layer numbers. Experimented on the 20 Newsgroups dataset with a pathological data distribution. Average 2 ranks were selected out of 16 ranks by our adaptive rank selection algorithm.