Federated Learning of Large Language Models with Parameter-Efficient Prompt Tuning and Adaptive Optimization

TL;DR

This paper tackles the challenge of federated fine-tuning for large language models by coupling parameter-efficient prompt tuning with adaptive optimization. It introduces FedPepTAO, which scores and selects layer-wise prompts to communicate, reducing both communication and computation without sacrificing accuracy, and applies a bidirectional adaptive optimization framework with per-device control variates to combat non-IID client drift. Extensive experiments across RoBERTaLARGE and multiple decoder-based LLMs demonstrate substantial gains in accuracy (up to 60.8%) and speed (up to 97.59% faster) over strong baselines, validating scalability to larger models and diverse tasks. The results suggest FedPepTAO as a practical approach for deploying FL with LLMs in real-world, data-siloed environments, while highlighting privacy considerations around shared prompts and directions for further mitigation.

Abstract

Federated learning (FL) is a promising paradigm to enable collaborative model training with decentralized data. However, the training process of Large Language Models (LLMs) generally incurs the update of significant parameters, which limits the applicability of FL techniques to tackle the LLMs in real scenarios. Prompt tuning can significantly reduce the number of parameters to update, but it either incurs performance degradation or low training efficiency. The straightforward utilization of prompt tuning in the FL often raises non-trivial communication costs and dramatically degrades performance. In addition, the decentralized data is generally non-Independent and Identically Distributed (non-IID), which brings client drift problems and thus poor performance. This paper proposes a Parameter-efficient prompt Tuning approach with Adaptive Optimization, i.e., FedPepTAO, to enable efficient and effective FL of LLMs. First, an efficient partial prompt tuning approach is proposed to improve performance and efficiency simultaneously. Second, a novel adaptive optimization method is developed to address the client drift problems on both the device and server sides to enhance performance further. Extensive experiments based on 10 datasets demonstrate the superb performance (up to 60.8\% in terms of accuracy) and efficiency (up to 97.59\% in terms of training time) of FedPepTAO compared with 9 baseline approaches. Our code is available at https://github.com/llm-eff/FedPepTAO.

Figures (9)

• Figure 1: The system model of FedPepTAO.
• Figure 2: The impact of server learning rate.
• Figure 3: The impact of various heterogeneity degrees.
• Figure 4: Evaluation of divers number of devices.
• Figure 5: Evaluation of divers layer selection strategies.