Communication-Efficient and Tensorized Federated Fine-Tuning of Large Language Models

2024-10-16

arXiv.org

AI Summary

To address the dual challenges of high communication overhead and severe data heterogeneity in federated fine-tuning of large language models (LLMs) over distributed private datasets, this paper proposes FedTT—a tensorized adapter architecture—and its enhanced variant, FedTT+. FedTT tensorizes parameter-efficient fine-tuning (PEFT) adapters and integrates them into a federated learning framework, enabling cross-device and cross-institutional deployment. FedTT+ further introduces an adaptive tensor factor freezing mechanism to significantly improve robustness under non-IID data distributions. The approach synergistically combines tensor decomposition, adapter injection (into encoder/decoder blocks), and federated optimization. Extensive evaluation on BERT and LLaMA demonstrates that FedTT and FedTT+ match or surpass state-of-the-art federated PEFT methods in accuracy while reducing communication costs by up to 10×. To our knowledge, this is the first work achieving highly efficient, robust, and low-overhead parameter-efficient federated fine-tuning of LLMs.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Parameter-efficient fine-tuning (PEFT) methods typically assume that Large Language Models (LLMs) are trained on data from a single device or client. However, real-world scenarios often require fine-tuning these models on private data distributed across multiple devices. Federated Learning (FL) offers an appealing solution by preserving user privacy, as sensitive data remains on local devices during training. Nonetheless, integrating PEFT methods into FL introduces two main challenges: communication overhead and data heterogeneity. In this paper, we introduce FedTT and FedTT+, methods for adapting LLMs by integrating tensorized adapters into client-side models' encoder/decoder blocks. FedTT is versatile and can be applied to both cross-silo FL and large-scale cross-device FL. FedTT+, an extension of FedTT tailored for cross-silo FL, enhances robustness against data heterogeneity by adaptively freezing portions of tensor factors, further reducing the number of trainable parameters. Experiments on BERT and LLaMA models demonstrate that our proposed methods successfully address data heterogeneity challenges and perform on par or even better than existing federated PEFT approaches while achieving up to 10$ imes$ reduction in communication cost.

Problem

• Reducing communication overhead in federated fine-tuning of LLMs
• Addressing data heterogeneity in distributed private datasets
• Enhancing parameter efficiency while maintaining model performance

Innovation

• Tensorized adapters for efficient federated fine-tuning
• Adaptive freezing to handle data heterogeneity
• Reduces communication costs by up to 10x