Robust Federated Finetuning of Foundation Models via Alternating Minimization of LoRA
Shuangyi Chen Yue Ju Hardik Dalal Zhongwen Zhu Ashish Khisti


[
Robust Federated Finetuning of Foundation Models 
 via Alternating Minimization of LoRA

















Shuangyi Chenyyy
Yue Juxxx
Hardik Dalalxxx
Zhongwen Zhuxxx
Ashish Khistiyyy


yyyECE Department, University of Toronto, Toronto, Canada
xxxEricsson-GAIA Montréal, Canada



Shuangyi Chenshuangyi.chen@mail.utoronto.ca
Ashish Khistiakhisti@ece.utoronto.ca




Machine Learning, ICML

0.3in
]









  


§ ABSTRACT 
Parameter-Efficient Fine-Tuning (PEFT) has risen as an innovative training strategy that updates only a select few model parameters, significantly lowering both computational and memory demands. PEFT also helps to decrease data transfer in federated learning settings, where communication depends on the size of updates. In this work, we explore the constraints of previous studies that integrate a well-known PEFT method named LoRA with federated fine-tuning, then introduce RoLoRA, a robust federated fine-tuning framework that utilizes an alternating minimization approach for LoRA, providing greater robustness against decreasing fine-tuning parameters and increasing data heterogeneity. Our results indicate that RoLoRA not only presents the communication benefits but also substantially enhances the robustness and effectiveness in multiple federated fine-tuning scenarios.




§ INTRODUCTION


The recent emergence of foundation models in various applications significantly changes the field of machine learning. Characterized by their broad adaptability and massive scale, these models require access to vast and diverse datasets to effectively learn across different tasks and domains. However, this presents a significant challenge: foundation models not only require large amounts of data, but also data of high quality.
Federated learning provides a promising solution to this issue. It enables the use of data from multiple sources while protecting the privacy of the data. By combining insights from different decentralized sources, federated learning allows for collaborative model training without exposing sensitive information. This method is especially beneficial for foundation models, as it can access a broad range of data while maintaining privacy. 

Recently, Parameter-Efficient Fine-Tuning (PEFT) has emerged as an innovative training strategy that updates only a small subset of model parameters, substantially reducing computational and memory demands. A notable method in this category is LoRA <cit.>, which utilizes low-rank matrices to approximate weight changes during fine-tuning. These matrices are integrated with pre-trained weights for inference, facilitating reduced data transfer in scenarios such as federated learning, where update size directly impacts communication efficiency. Many works integrate LoRA into federated setting. For example, FedPETuning <cit.> compared various PEFT methods in a federated setting. SLoRA <cit.>, a hybrid approach that combines sparse fine-tuning with LoRA, is introduced to tackle data heterogeneity in federated settings. Furthermore, FS-LLM <cit.> is presented, which is a framework for fine-tuning LLMs in federated environments. However, these studies typically apply the FedAVG algorithm directly to LoRA modules, overlooking the interference introduced by this aggregation approach. With this consideration, Sun et al. designs a federated finetuning framework named FFA-LoRA <cit.> based on LoRA by freezing down-projection matrix 𝐀 for all the clients and only updating up-projection matrix 𝐁. Furthermore, they apply DP-SGD to preserve privacy. Using sufficient number of finetuning parameters, FFA-LoRA with a larger learning rate achieves performance comparable to FedAVG for LoRA modules while halving the communication costs. However, we observe that with fewer fine-tuning parameters, FFA-LoRA is less robust than FedAVG for LoRA modules, primarily due to its limited expressiveness stemming from the restricted number of trainable parameters. 
Another common issue in federated learning is data heterogeneity among clients. To address this, we drew inspiration from the personalized federated framework FedRep <cit.>, which alternates between updating clients' representation and head. This approach highlights the importance of learning a robust low-rank representation and demonstrates superior convergence speed compared to simultaneously updating both representation and head.
Therefore, we propose a robust federated fine-tuning framework, RoLoRA, based on alternating minimization of LoRA. Empirical evidence demonstrates that RoLoRA is more robust against Decreasing Fine-tuning Parameters and Increasing Data Heterogeneity, while still halving communication costs, similar to FFA-LoRA.


  
Related Work We provide a summary of the literature on PEFT, Variants of LoRA, PEFT in Federated Setting, and FL with data heterogeneity in  Appendix <ref>.















§ PRELIMINARIES



 §.§ Low-Rank Adaptation: LoRA
 
Low-Rank Adaptation (LoRA) <cit.> fine-tunes large language models efficiently by maintaining the original model weights fixed and adding small, trainable matrices in each layer. These matrices perform low-rank decompositions of updates, reducing the number of trainable parameters. This approach is based on the finding that updates to model weights during task-specific tuning are usually of low rank, which allows for fewer parameters to be adjusted. For example, for a pre-trained weight matrix 𝐖_0 ∈ℝ^d× d, the update is a low-rank product 𝐁𝐀, where 𝐀∈ℝ^r× d and 𝐁∈ℝ^d× r, with r << d. Only 𝐀 and 𝐁 are trainable, allowing 𝐖 = 𝐖_0 + α𝐁𝐀, with α adjusting the update's impact.
Applying LoRA in a federated setting is a practical choice. By using LoRA adapters, clients can fine-tune foundation models efficiently with limited resources. Since only these specific matrices need to be transmitted to a central server, this approach significantly reduces communication costs. This makes LoRA an advantageous solution for enhancing model performance in collaborative scenario comparing to full parameter finetuning in the federated setting. 


 §.§ FedAVG of LoRA Introduces Interference

Integrating LoRA within a federated setting presents challenges. In such a setup, each of the N clients is provided with the pretrained model weights 𝐖_0, which remain fixed during finetuning. Clients are required only to send the updated matrices 𝐁_i and 𝐀_i to a central server for aggregation. While most current studies, such as SLoRA<cit.> and FedPETuning<cit.>, commonly apply FedAVG directly to these matrices as shown in (<ref>), this approach might not be optimal. The precise update for each client’s model, Δ𝐖_i, should be calculated as the product of the low-rank matrices 𝐀_i and 𝐁_i. Consequently, aggregation on the individual matrices introduces interference. 

    1/N∑_i=1^N  Δ𝐖_i = 1/N (𝐁_1𝐀_1+𝐁_2𝐀_2+...+ 𝐁_N𝐀_N) 
    ≠1/N (𝐁_1+𝐁_2+...+𝐁_𝐍) 1/N (𝐀_1+𝐀_2+...+𝐀_𝐍)



 §.§ FedRep: Common Representation via Alternating Minimization
 
A common challenge in federated learning is data heterogeneity among clients. FedRep <cit.> addresses this by finding a common representation, which is effectively achieved via alternatively updating clients' representation and head, aggregating only representation while keeping the head diverse. The algorithm demonstrates the necessity of learning a robust low-rank representation. Additionally, the alternating optimization has shown superior convergence speed compared to approaches that simultaneously update both representation and head. We observe a structured similarity between the LoRA adapter and the representation-head structure. Specifically, we consider the down-projection matrix 𝐀 in each LoRA adapter as the low-rank representation for the features of the intermediate layers. We hypothesize that learning a robust low rank representation (down-projection matrix) is also advantageous for the intermediate features when the clients has heterogeneous inputs. However, since the LoRA adapters are cascaded in the model unlike single representation-head structure in model considered in FedRep, keeping up-projection matrix 𝐁 diverse may not be favorable for convergence. 

With these considerations, we propose robust federated fine-tuning framework based on alternating minimization of LoRA (RoLoRA).

























§ OUR FRAMEWORK
 
We describe the framework design of RoLoRA and discuss its practical advantages. 


  
Alternating Minimization and Corresponding Aggregation
Motivated by the observations discussed in Section <ref> and <ref>, we propose applying alternating minimization to the local fine-tuning of each client in a setting with N clients. Unlike the approach in FFA-LoRA, where 𝐀 is consistently frozen, we suggest a alternating update strategy. There are alternating odd and even communication rounds designated for updating, aggregating 𝐀 and 𝐁, respectively. 
    In the odd comm. round:    1/N∑_i=1^N  Δ𝐖_i^2t+1
       = 1/N (𝐁_1^t+1𝐀_1^t+𝐁_2^t+1𝐀_2^t+...+ 𝐁_N^t+1𝐀_N^t) 
       = 1/N (𝐁_1^t+1+𝐁_2^t+1+...+ 𝐁_N^t+1)𝐀^t  
       In the even comm. round:  1/N∑_i=1^N  Δ𝐖_i^2t+2
       = 1/N (𝐁_1^t+1𝐀_1^t+1+𝐁_2^t+1𝐀_2^t+1+...+ 𝐁_N^t+1𝐀_N^t+1) 
       = 1/N𝐁^t+1(𝐀_1^t+1+𝐀_2^t+1+...+ 𝐀_N^t+1)

 
In the odd communication round, all clients freeze 𝐀^t and update 𝐁^t. The central server then aggregates these updates to compute 𝐁^t+1 = 1/N∑_i=1^N𝐁^t+1_i and distributes 𝐁^t+1 back to the clients. In the subsequent communication round, clients freeze 𝐁^t+1 and update 𝐀^t. The server aggregates these to obtain 𝐀^t+1 = 1/N∑_i=1^N𝐀^t+1_i and returns 𝐀^t+1 to the clients. 
It is important to note that in round 2t+1, the frozen 𝐀_i^t are identical across all clients, as they are synchronized with 𝐀^t from the central server at the beginning of the round. This strategy ensures that the update and aggregation method introduces no interference, as demonstrated in (<ref>) and (<ref>).











  
Computation and Communication Cost
The parameter-freezing nature of RoLoRA enhances computational and communication efficiency. In each communication round, the number of trainable parameters in the model is effectively halved compared to FedAVG with LoRA. The only additional cost for RoLoRA compared to FFA-LoRA is the alternating freezing of the corresponding parameters. We remark this additional cost is negligible because it is applied to the clients' models and can be executed concurrently during the server's aggregation.


§ EXPERIMENTS

We evaluate the performance of RoLoRA in various federated settings. We use NVIDIA GeForce RTX 4090 or NVIDIA A40 for all the experiments.




  
Baselines Considering cross-silo federated setting where the number of clients is relatively small and all clients will participate in each round, we will explore the following three methods based on FedAVG.

    
  * LoRA means LoRA adapter and its finetuning algorithm are directly applied to local finetuning of clients in the federated system. Specifically, in iteration t, the server receives 𝐀_i^t and 𝐁_i^t from client i and aggregates by  𝐀^t = 𝖠𝗏𝗀(𝐀_i^t) and 𝐁^t = 𝖠𝗏𝗀(𝐁_i^t). 
    
  * LoRA-FFA <cit.> is a baseline that enable the clients to finetune 𝐁 and keep 𝐀 frozen locally. Thus, in iteration t, the server aggregates by 𝐁^t = 𝖠𝗏𝗀(𝐁_i^t). 
    
  * RoLoRA enables clients to alternate updating 𝐀 and 𝐁 as described in Section <ref>.  





  
Model and Datasets. We take the pre-trained RoBERTa-Large (355M) <cit.> models from the HuggingFace Transformers library. and evaluate the performance of three federated finetuning methods on 5 datasets (SST-2, QNLI, MNLI, QQP, RTE) from the GLUE <cit.>. Due to the limitation of the unpublished test set in GLUE, we follow the previous studies <cit.> and use the original validation set as the new test set and split a part of the training set as the validation set. 



  
Implementation.We implement all the methods based on FederatedScope-LLM <cit.>. To make a fair comparison, for each dataset, we obtain the best performance on test set and report the average over five seeds. Specifically, the learning rate is chosen from the set {5e-4, 1e-3, 2e-3, 5e-3, 1e-2, 2e-2, 5e-2, 1e-1, 2e-1}. Other hyper-parameters for experiments are specified in Table <ref> in Appendix <ref>.
















































    
















































































































  
Effect of Number of Finetuning Parameters
In Figure <ref>, we compare three methods across five GLUE datasets. We apply LoRA to every weight matrix of the selected layers, given different budgets of LoRA parameters. For each dataset, we experiment with three budgets, ranging from high to low. The corresponding layer sets, 𝒫_1, 𝒫_2, 𝒫_3, are detailed in Table <ref> in Appendix <ref>.

The figures indicates that with sufficient number of finetuning parameters, the three methods can achieve comparable best accuracy; as the number of LoRA parameters is reduced, the performance of the three methods deteriorates to varying degrees. However, RoLoRA, which achieves performance comparable to LoRA, demonstrates greater robustness compared to FFA-LoRA, especially under conditions of limited fine-tuning parameters. It is important to note that with the same finetuning parameters, the communication cost of RoLoRA and FFA-LoRA is always half of that of LoRA due to their parameter freezing nature. This implies that RoLoRA not only sustains its performance but also enhances communication efficiency. We expand the middle set of data of each of Figure <ref>, corresponding to 𝒫_2, and show the details of the performance of three methods in Table <ref>. 





















































































  
Effect of Data Heterogeneity

In this section, we study the effect of data heterogeneity. The layer set with LoRA adapters in Table <ref> is 𝒫_2 as in Table <ref>. In Table <ref>, we increased the number of clients from 3 to 20, and then to 50, ensuring that there is no overlap in the training samples each client can access. Consequently, each client receives a smaller fraction of the total dataset, leading to a rise in data heterogeneity among the clients. We observe that as the data heterogeneity increases, while maintaining the same number of fine-tuning samples, the performance of the LoRA method significantly deteriorates for most datasets. In contrast, RoLoRA maintains its accuracy levels. The performance of FFA-LoRA also declines, attributed to the limited expressiveness of the random initialization of 𝐀 for clients' heterogeneous data. Notably, RoLoRA achieves this accuracy while incurring only half the communication costs associated with LoRA. Figure <ref> in Appendix <ref> illustrates the dynamics during fine-tuning for three methods, highlighting that the convergence speed of RoLoRA is substantially better than that of the other two methods.



  
Align Communication Cost for Three Methods

In Table <ref>, we conduct a comparison of three methods under the constraint of identical communication costs under the assumption that the number of clients is small. To align the communication costs across these methods, two approaches are considered. The first approach involves doubling the rank of FFA-LoRA and RoLoRA, with the outcomes detailed in Table <ref>. The second approach requires doubling the number of layers equipped with LoRA adapters. In the results presented in Table <ref>, the latter strategy is employed.  Specifically, for both FFA-LoRA and RoLoRA, we adjust the communication costs by doubling the number of layers equipped with LoRA adapters, compared to the baseline LoRA method, where the layer set 𝒫_3 are attached with adapters, thereby standardizing the size of the transmitted messages. Table <ref> demonstrates that when operating within a constrained communication cost budget, the performance of RoLoRA consistently surpasses that of the other two methods.



More experimental results with different models and settings are provided in Appendix <ref>.


§ CONCLUSION

In this work, we introduce RoLoRA, a robust federated fine-tuning framework using alternating minimization for LoRA. RoLoRA improves robustness against reduced fine-tuning parameters and increased data heterogeneity. Our results show that RoLoRA enhances communication efficiency, robustness, and effectiveness in various federated fine-tuning settings.

icml2024









§ RELATED WORKS
 


 §.§ Parameter Efficient Fine Tuning (PEFT): LoRA and Its Variants

As the size of large language models (LLMs) continues to increase, it is computationally expensive and time-consuming to finetune to the full model. Parameter efficient finetuning (PEFT) allows for updates to a smaller subset of parameters, significantly reducing the computational and memory requirements. One of the most well-known methods is LoRA<cit.>.
LoRA uses low-rank matrices to approximate changes in weights during fine-tuning, allowing them to be integrated with pre-trained weights before inference. Based on LoRA, many PEFT methods are developed. For example, Zhang et al. <cit.> designs AdaLoRA by using SVD decomposition and pruning less significant singular values for more efficient updates. VeRA <cit.> is proposed to further reduce the number of trainable parameters during finetuning by using a single pair of low-rank matrices shared across all layers and learning small scaling vectors. Zhang et al. <cit.> proposes a memory-efficient fine-tuning method named LoRA-FA which keeps the projection-down weight of 𝐀 fixed and updates the projection-up weight of 𝐁 during finetuning. Hayou et al. <cit.> enhance LoRA by assigning different learning rates to 𝐀 and 𝐁, theoretically confirming that the optimal approach requires a higher learning rate for 𝐁 than for 𝐀. Liu et al. analyze magnitude and directional updates in LoRA versus full parameter fine-tuning and introduce DoRA<cit.>, which decomposes pre-trained weights for fine-tuning and applies LoRA for directional updates. A quantized version of LoRA named QLoRA<cit.> is introduced. Building upon that, Li et al. develops LoftQ <cit.> for a better initialization for quantized training.


 §.§ PEFT in Federated Setting


PEFT adjusts only a few lightweight or a small portion of the total parameters for specific tasks, keeping most foundational model parameters unchanged. This feature can help reduce data transfer in federated learning, where communication depends on the size of updates. Zhang et al. <cit.> compares multiple PEFT methods in federated setting, including Adapter<cit.>, LoRA<cit.>, Prompt tuning<cit.> and Bit-Fit<cit.>. SLoRA<cit.>, which combines sparse finetuning and LoRA, is proposed by Babakniya et al. to address the data heterogeneity in federated setting. 

Sun et al. designs a federated finetuning framework named FFA-LoRA based on LoRA <cit.> by freezing matrix 𝐀 for all the clients and only updating matrix 𝐁. Furthermore, they apply DP-SGD to preserve privacy. FS-LLM <cit.>, a framework for finetuning LLM, is introduced.


 §.§ FL with Data Heterogeneity



  
FL with a Common Representation
FL with a common representation aims to address the challenges of data heterogeneity in FL. Those FL methods learn a shared global representation while allowing each client to have its own personalized partial model. Works include FedRep<cit.>, which learns a global low-dimensional representation and personalized head for each client, FedCR <cit.>, which introduces a regularizer to encourage learning a shared representation, and FedPAC <cit.>, which performs class-wise feature alignment. Other methods like FedBABU <cit.> and FedRoD <cit.> also aim to learn a shared representation across clients. Although we focus on a common model in the federated setting in this work, we got inspired by training algorithm introduced by FedRep <cit.> to learn a low-rank representation for the intermediate features. We discuss the similarity and difference between the LoRA adapter and representation-head structure. 



§ EXPERIMENTS
 


 §.§ Setup
 
We show the hyper-parameter configurations for each dataset in Table <ref>.



In Table <ref>, we include the details about layers attached with LoRA adapters for different budget of finetuning parameters, for each dataset.


 §.§ More Results
 





















 


  §.§.§ Communication Cost

In Table <ref>, we show the uplink communication cost for three methods for the layer set 𝒫_2 using rank=1.




  §.§.§ Finetuning Dynamics of the Setup with Severe Data Heterogeneity
 
In Figure <ref>, we show the convergence of three methods under severe data heterogeneity with 50 clients. RoLoRA demonstrates superior convergence speed compared to the other two methods.




  §.§.§ DeBERTa-XLarge Results
 
In Table <ref>, we show the results with DeBERTa-XLarge (900M). For both FFA-LoRA and RoLoRA, we modify the communication costs by equipping twice as many layers with LoRA adapters compared to the standard LoRA method. So the communication costs are aligned for three methods.




  §.§.§ Performance when QLoRA is Applied
 In Table <ref>, we quantize the frozen pre-trained weights to 8 bit and 4 bit for each client, and apply QLoRA<cit.>. The relative accuracy is computed as 𝖠𝖼𝖼_𝖥𝖯-𝖠𝗏𝗀(𝖠𝖼𝖼_8b+𝖠𝖼𝖼_4b)/𝖠𝖼𝖼_𝖥𝖯. We use sufficient finetuning parameters and the selected layer set is 𝒫_1 to study the effect of the quantized foundation model in the federated setting.





