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

2024-10-12

AI Summary

In federated learning, LoRA fine-tuning suffers from update bias due to FedAvg’s averaging of low-rank adapters, leading to inaccurate global model aggregation. To address this, we propose a residual error compensation mechanism: under frozen pretrained weights, we explicitly model and cumulatively correct the aggregation residuals of LoRA adapters, enabling strictly lossless weight aggregation. This is the first method to support *exact* (non-approximate) gradient/weight aggregation for LoRA in federated settings, balancing communication efficiency and computational overhead. Experiments across arithmetic reasoning, commonsense reasoning, natural language understanding, and generation tasks demonstrate significant improvements over existing SOTA approaches. Results confirm that update bias substantially degrades performance and validate our method’s strong robustness and generalization capability across diverse NLP benchmarks.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Distributed Machine Learning & Federated Learning
• Multiagent Systems: Adversarial Agents

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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 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. Our code is publicly available at https://github.com/RaghavSinghal10/fedex-lora.

Problem

• Enables exact aggregation in federated learning.
• Improves efficiency in fine-tuning foundation models.
• Addresses challenges in distributed data environments.

Innovation

• Exact updates with minimal overhead
• Residual error term in LoRA
• Efficient federated fine-tuning