Ravan: Multi-Head Low-Rank Adaptation for Federated Fine-Tuning

Thank you for the thorough assessment of our paper and valuable feedback. We hope to comprehensively respond to each of your concerns.

W1. Results are reported on simulated data, where the heterogeneity has been introduced artificially. It is unclear if the method can capture heterogeneity in real datasets. It would be useful to report performance on alternative datasets used in FL research e.g. FedJAX (https://github.com/google/fedjax) contains prepackaged datasets such as EMNIST-62, Stack Overflow

Thank you for the suggestion; based on this comment we have included results on federated EMNIST-62 where the dataset has a natural non-I.I.D. partition based on the writer responsible for specific handwritten characters. We use the ViT-B/16 model and the hyperaparameter setup described in the lower halves of Tables 2 and 3 in Section 4 of our paper. We train 3 different settings in which we randomly sample a total of 20 clients, 50 clients, or 100 clients for model fine-tuning and where each client corresponds to a specific writer and their dataset corresponds to their handwritten characters.

FEMNIST-62 with 20 Clients -

FEMNIST-62 with 50 Clients -

FEMNIST-62 with 100 Clients -

As is evident, even when using a dataset with a natural non-I.I.D. partition, Ravan outperforms related federated PEFT methods. Thus, the method is not just applicable in simulated non-I.I.D. datasets but also practically useful in real-world settings. We will include this experiment in our final paper.

Q1. Could you provide more details about how the clients choose # of heads based on their computational capacity? Also, it would be useful to show the distribution of # of heads chosen across clients for each of the 3 distributions.

For the sake of space, the discussion on how clients choose the number of heads for fine-tuning was included in the Appendix as opposed to the main text of the paper. The distribution of number of heads chosen across clients for the 3 distributions can be found in Figure 5 of the Appendix. We define $N_{\max}$ as the largest number of trainable parameters that any client can afford. Each client $c$ is constrained to a trainable parameter budget $N_c \in \\{\frac{1}{4}, \frac{1}{2}, \frac{3}{4}, 1\\} N_{\max}$. A client with budget $N_c$, fine-tunes only $\bigl\lfloor \left(N_c \\, \/ \\, N_{\max}\right) \cdot H \bigr\rfloor$ heads and leaves the remaining heads frozen (e.g. for $N_c = \frac{1}{4} N_{\max}$ the client trains one quarter of the heads).

Q2. It would be good to understand what happens when $\sqrt{N*h}$ exceeds d. Why does the effective rank of each head become smaller in this setup? Could you provide further insight?

An important thing to note for Figure 4 is that while we increase the number of heads, we are still keeping the number of trainable parameters constant. In order to do so, the size of each individual head must decrease, otherwise the number of trainable parameters would also increase. As a consequence, the effective rank of each individual head goes down. This is fine as long as the effective rank of the overall sum continues to increase. However, beyond a certain number of heads, the effective rank of the overall sum equals the embedding dimension $d$, which is the maximum possible rank of the update approximation, and no longer increases. At this point, as each individual head becomes less expressive, model performance begins to degrade. This phenomenon can be observed in our results figure.

Thank you for the detailed critique of our work and the recognition of its value in the field of federated fine-tuning. We aim to thoughtfully address the specific points you have raised.

W1. The claim "To our knowledge, RAVAN is the first federated PEFT technique that simultaneously maintains parameter-efficient computation and communication and addresses both data and computational heterogeneity" is too strong. The idea has been around for a long time, with works like FedPara and others.

We appreciate the feedback on the framing of our contributions and will revise this sentence in the final draft of our paper to better reflect the benefits of Ravan in the context of current research. The following will be included in the paper should it be accepted "Extending prior efforts such as FedPara, Ravan introduces an integrated framework that maintains parameter‑efficient computation and communication and demonstrates robustness across diverse data and computational heterogeneity in federated settings."

W2. For part 3.1: why does Gram-Schmidt perform better? Is there a theoretical justification?

Section 3.1 discusses various strategies for initializing the Ravan parameters. In particular, we find that the Gram-Schmidt initialization and random normal initialization are effective initializations for the $**B**_i$ and $**A**_i$ parameters. The random normal initialization yields orthogonality for the $**B**_i$ and $**A**_i$ parameters in expectation, while Gram-Schmidt ensures that the $**B**_i$ and $**A**_i$ parameters are orthogonal deterministically. This orthogonality is important because the column space of $\sum s_i**B**_i**H**_i**A**_i$ is a subspace of the span of the $**B**_i$'s and the row space of $\sum s_i**B**_i**H**_i**A**_i$ is a subspace of the span of the $**A**_i$'s. Therefore, if we constrain all the $**B**_i$'s and $**A**_i$'s to share the same subspace, the rank of $\sum s_i**B**_i**H**_i**A**_i$ collapses to that of a single head. This is the case for both the constant initialization, where $**B**_i = **B**_j$ and $**A**_i = **A**_j \ \forall i,j$, and in the shared subspace initialization, where $**B**_i = **MR**_i, **A**_i = **R**_i**N**$ for normally distributed $**M**, **N**$ and invertible $**R**_i \ \forall i$. For these initializations, we would not benefit from using multiple heads at all since the effective rank of the update approximation would not actually improve. Section 4.3 demonstrates that initializations that do improve the rank of the update approximation (random normal, Gram-Schmidt) outperform methods that do not improve the rank of the update approximation (constant, shared subspace).

Further empirical evidence of the improvement to the rank of the update approximation can be demonstrated by the following experiment. For CIFAR-100, we calculated the rank of the update approximation $\sum s_i**B**_i**H**_i**A**_i$ for all 4 initialization schemes (constant, shared subspace, random normal, Gram-Schmidt) after communication round 1 and communication round 50. For the sake of calculating the rank, any singular value less than $1 \times 10^{-6}$ is considered 0. We follow the same hyperparameter setup in Tables 2 and 3 by using 4 heads of rank 156 each.

After communication round 1 -

After communication round 50 -

These results establish that Ravan achieves the upper bound theoretical rank of 624 ($4 \times 156$) throughout the training procedure when using initializations that result in orthogonal $**B**_i$ and $**A**_i$.

W3. It would also be interesting to compare these results with a pre-trained initialization on the same domain using a random dataset.

This is an exciting idea and a potential alternative initialization scheme to the ones we include in the paper. We have benchmarked this initialization in which we pretrain the $**B**_i$ and $**A**_i$ parameters prior to the FL fine-tuning procedure. The $**B**_i$ and $**A**_i$ parameters are subsequently frozen at this pretrained initialization. The following are results on CIFAR-100 and SVHN in which the $**B**_i$ and $**A**_i$ parameters are pretrained on random subsets of CIFAR-100 and SVHN, respectively. The hyperparameter settings follow the ones listed in Table 4.

Results for CIFAR-100 -

Results for SVHN -

The results of this initialization are mixed with improvements on our proposed initializations in CIFAR-100 but a slight decrease in performance on SVHN (though we admit that testing is limited given the time constraints of the rebuttal period). While this direction is certainly interesting, and a potential direction for future work, we would like to bring up that this initialization would likely require pretraining the $**B**_i$ and $**A**_i$ parameters at the server since they are quite a bit larger than the $**H**_i$ parameters and might be difficult to train given device resource constraints. This would require a public dataset at the server that could be used to pretrain the $**B**_i$ and $**A**_i$ parameters. Thus, there are some limitations that might make this initialization more difficult in practice. The initializations that we propose are light-weight, effective, and data-independent. That being said, we do believe that the initialization you propose is an intriguing direction for future study.

W4. The choice of 50 local epochs for client training is very problematic for comparing results with other works in the literature. In federated learning, it is common to see clients training for 1, 5, 10, or 20 epochs.

We believe there is a misunderstanding here -- we train for 50 local iterations and not 50 local epochs. This means that for each client, we only train on 50 mini-batches and not 50 entire traversals of the client's training dataset. 50 local iterations, depending on the dataset size and batch size, is approximately 1 local epoch (client dataset size / batch size = local iterations per epoch). We use local iterations instead of local epochs since clients may have different amounts of training data, so using local iterations means that each client performs exactly the same number of forward-backward passes. We believe that this is a fair amount of compute for resource-constrained devices and does not differ from most modern FL literature in a significant way. We will highlight the difference between local iterations and local epochs in the final draft of our paper to avoid similar confusion in the future. However, to demonstrate the robustness of our method, we include the results for CIFAR-100 and SVHN using 1 or 2 local epochs below.

Results for CIFAR-100 with 1 Local Epoch -

Results for CIFAR-100 with 2 Local Epochs -

Results for SVHN with 1 Local Epoch -

Results for SVHN with 2 Local Epochs -

W5. The paper mentions that the "scaling factors" are learned, but it is not clear how this learning process occurs. This aspect needs to be clarified for better understanding.

Effectively, the scaling factors $s_i$ are trainable parameters just like any other trainable parameter in the network. The only difference is that they are scalars and not multi-dimensional tensors. However, they receive gradients, just like the $**H**_i$ parameters, and are trained with these gradient updates. The advantage to keeping these $s_i$ factors trainable is that the network naturally learns which heads are more significant throughout the fine-tuning procedure and as a consequence adjusts the value of $s_i$, in essence learning the scaling factors most beneficial for model prediction.

Thank you for the meaningful feedback and the acknowledgement of Ravan's strengths and practical utility. We would like to address your individual concerns.

W1. A more rigorous theoretical justification would be appreciated to formally connect the spectral properties of non-I.I.D. updates to the necessity of a higher-rank adaptation space.

Thank you for bringing up this discussion; we believe this is a valid point about the theoretical rigor of our proposed method. Unfortunately, there exists limited theory about the spectral properties of LoRA optimization, especially in distributed or federated settings. As such, our work closely aligns with previous works in federated fine-tuning that focus primarily on the empirical advantages of methodological innovations ([1], [2], [3]). That being said, there are a few prior works ([4], [5], [6]) that tie the convergence of LoRA directly to the rank and suggest that when the actual desired update is high rank, LoRA may underperform or even fail to converge. This is particularly important given our observation that the true update is high rank in non-I.I.D. settings. In this context, we believe there are two critical pieces of evidence that suggest that the rank of the update approximation contributes to the improvement in performance in federated settings.

• Increasing the number of heads, and subsequently the rank of the update approximation, improves the performance of the method only until the rank of the update approximation stops increasing. In Section 4.3, we include an ablation that demonstrates the impact of changing the number of heads while leaving the number of trainable parameters constant. We find that increasing the number of heads is beneficial as long as the rank of the update approximation remains less than the embedding dimension, which is the maximum possible rank of the update approximation. However, once the rank of the update approximation equals the embedding dimension and stops increasing, these performance gains vanish. Since we hold the number of trainable parameters constant in this experiment, this conveys that the rank of the update approximation is associated directly with the performance of the method.
• Initialization schemes that do not increase the effective rank of the update approximation perform worse on all benchmarks. In Section 3.1, we propose two initialization benchmarks that do not increase the effective rank of the update approximation. In the constant initialization $**B**_i = **B**_j$ and $**A**_i = **A**_j \ \forall i,j$, and in the shared subspace initialization $**B**_i = **MR**_i, **A**_i = **R**_i**N**$ for normally distributed $**M**, **N**$ and invertible $**R**_i \ \forall i$. We highlight that the constant initialization and shared subspace initialization underperform relative to the random normal initialization and Gram-Schmidt initialization. Despite using multiple heads, the constant and shared subspace initializations have an update approximation that is the same rank as a single head. In contrast, the random normal initialization and Gram-Schmidt initialization have an effective rank that is (number of heads * rank of each head). Since each initialization uses the same number of trainable parameters, the same rank for each head, and the same number of heads, the only variable that changes is the effective rank of the entire update approximation which qualifies the importance of increasing the rank of the update approximation.

While we agree that our work presents limited theoretical insights into this behavior, we believe these experiments and prior theoretical work combined present convincing evidence that the improved rank of Ravan's update approximation is a critical component of the method's performance.

Q1. While Ravan offers a high theoretical rank, does the optimization process for non-I.I.D. updates actually utilize this full capacity?

Yes, the optimization process actually utilizes this full capacity as demonstrated by the following experiment. For CIFAR-100, we calculated the rank of the update approximation $\sum s_i**B**_i**H**_i**A**_i$ for all 4 initialization schemes (constant, shared subspace, random normal, Gram-Schmidt) after communication round 1 and communication round 50. For the sake of calculating the rank, any singular value less than $1 \times 10^{-6}$ is considered 0. We follow the same hyperparameter setup in Tables 2 and 3 by using 4 heads of rank 156 each.

After communication round 1 -

After communication round 50 -

These results establish that Ravan achieves the upper bound theoretical rank of 624 ($4 \times 156$) throughout the training procedure when using initializations that result in orthogonal $**B**_i$ and $**A**_i$. It also lends further credence to the idea that the differentiating factor in performance between the various initialization schemes is a large difference in effective rank. We have included the performance for the different initialization schemes on CIFAR-100 here for convenience.

Results for CIFAR-100 -

[1] "FedEx-LoRA: Exact Aggregation for Federated and Efficient Fine-Tuning of Foundation Models" Singhal et al.

[2] "FLoRA: Federated Fine-Tuning Large Language Models with Heterogeneous Low-Rank Adaptations" Wang et al.

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

[4] "Randomized Asymmetric Chain of LoRA: The First Meaningful Theoretical Framework for Low-Rank Adaptation" Malinovsky et al.

[5] "GeLoRA: Geometric Adaptive Ranks for Efficient LoRA Fine-Tuning" Ed-dib et al.

[6] "LoRA Training Provably Converges to a Low-Rank Global Minimum or It Fails Loudly (But it Probably Won't Fail)" Kim et al.