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

Raghav Singhal, Kaustubh Ponkshe, Praneeth Vepakomma

Track: Main paper track (up to 5 pages excluding references and appendix)

Keywords: Fine-Tuning, Parameter-Efficient Fine-Tuning, PEFT, LoRA, Federated Learning

TL;DR: Our method achieves exact parameter aggregation in federated fine-tuning using LoRA, resulting in better performance.

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 **Fed**erated **Ex**act **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 models across arithmetic reasoning, commonsense reasoning, natural language understanding and natural language generation 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.

Anonymization: This submission has been anonymized for double-blind review via the removal of identifying information such as names, affiliations, and identifying URLs.

Submission Number: 50