FedEx-LoRA: Exact Aggregation for Federated and Efficient Fine-Tuning of Foundation Models

Raghav Singhal*1, Kaustubh Ponkshe*1, Praneeth Vepakomma1,2
1 Mohamed bin Zayed University of Artificial Intelligence, UAE
2 Massachusetts Institute of Technology, USA
* Indicates Equal Contribution
A version of our paper is accepted at NeurIPS 2024 Workshop on Fine-Tuning in Modern Machine Learning: Principles and Scalability (FITML)
Code arXiv

Abstract

Low-Rank Adaptation (LoRA) is a popular technique for efficient fine-tuning of foundation models. However, applying LoRA in federated learning environments, where data is distributed across multiple clients, presents unique challenges. Existing methods rely on traditional federated averaging of LoRA adapters, resulting in inexact updates. To address this, we propose Federated Exact LoRA, or FedEx-LoRA, which adds a residual error term to the pretrained frozen weight matrix. Our approach achieves exact updates with minimal computational and communication overhead, preserving LoRA’s efficiency. We evaluate the method on various Natural Language Understanding (NLU) and Natural Language Generation (NLG) tasks, showing consistent performance gains over state-of-the-art methods across multiple settings. Through extensive analysis, we quantify that the deviations in updates from the ideal solution are significant, highlighting the need for exact aggregation. Our method’s simplicity, efficiency, and broad applicability position it as a promising solution for accurate and effective federated fine-tuning of foundation models.

Why FedEx-LoRA?

FedIT method
(a) FedIT (typical federated LoRA approaches)
FedEx-LoRA method
(b) FedEx-LoRA (our proposed method)
Comparison of federated LoRA methods: (a) FedIT averages the individual client low-rank adapters Ai and Bi, resulting in inexact updates. (b) FedEx-LoRA sends the error residual ΔWres along with the individual adapters Ai and Bi, which is added to the pretrained weight matrix W0, ensuring exact aggregation. Clients transmit low-rank adapters Ai and Bi in both methods.

Comparisons with Baselines

Centralized LoRA (in grey) sets the benchmark skyline for its federated versions. Best results among federated methods (in blue) are highlighted in bold for each setting. FedEx-LoRA consistently outperforms existing methods across various tasks and settings.

Deviation/Divergence Analysis

We measure the scaled Frobenius norm of the divergence between the updates produced by FedAvg and the ideal LoRA updates, revealing several notable patterns. We observe that the divergence/deviation - (1) decreases as the model depth increases, (2) grows with a higher number of local epochs, (3) is more pronounced in the query (Q) matrices compared to the value (V) matrices, (4) consistently decreases as the number of aggregation rounds increases, both for the first-layer query matrix and for the average of the query and value matrices across all layers.

BibTeX

@misc{singhal2024exactaggregationfederatedefficient,
        title={Exact Aggregation for Federated and Efficient Fine-Tuning of Foundation Models}, 
        author={Raghav Singhal and Kaustubh Ponkshe and Praneeth Vepakomma},
        year={2024},
        eprint={2410.09432},
        archivePrefix={arXiv},
        primaryClass={cs.DC},
        url={https://arxiv.org/abs/2410.09432}, 
  }