Truncate without Fear: Module Aggregation and Redistribution in Federated Low-Rank Adaptation

Zhijie Chen, Yuxing Liu, Arindam Banerjee

Keywords: Federated Learning, Low-rank adaptation, model merging

Abstract:

While low-rank adaptations (LoRA) have shown promise as an efficient fine-tuning technique in federated learning (FL) to reduce communication complexity, the practical application requires careful attention to the challenges posed by the aggregation schemes on client modules. In this paper, we introduce TFLoRA, which directly optimizes over the adapter weights $W = BA^\top$, and redistributes the LoRA modules using the updated adapter weights. Our theoretical analysis shows the truncation error introduced during the redistribution step is mild and TFLoRA
achieves an $O(1/\sqrt{T})$ convergence rate. Compared to the existing methods, TFLoRA supports a wide range of optimizers on the server side and maintain the advantages in low communication overhead. We show empirical evidence that TFLoRA achieves better performance than the state-of-the-art federated LoRA mechanisms on various benchmarks including image/text classification and commonsense inference. Additionally, TFLoRA is demonstrated to be more favorable as the number of clients increases and with non-i.i.d client data distributions.

Submission Number: 33