FedRPCA: Enhancing Federated LoRA Aggregation Using Robust PCA
This work addresses data heterogeneity in federated learning for resource-constrained clients, but it is incremental as it builds on existing model merging techniques like Task Arithmetic.
The paper tackled the problem of data heterogeneity in federated learning with LoRA fine-tuning, which causes slow convergence and suboptimal accuracy, and proposed FedRPCA to decompose client updates into common and client-specific components, achieving higher final accuracy and faster convergence across vision and language tasks.
LoRA has emerged as one of the most promising fine-tuning techniques, especially for federated learning (FL), since it significantly reduces communication and computation costs at resource-constrained clients. However, data heterogeneity remains a significant challenge for LoRA-based FL, and the conventional aggregation strategy based on FedAvg suffers from slow convergence and suboptimal accuracy. Motivated by recent advances in model merging, particularly Task Arithmetic, we explore the idea of aggregating client LoRA parameters using scaled averaging. We first observe that a naive application of Task Arithmetic is ineffective due to the high cosine similarity between client updates, indicating significant common knowledge in the updates across clients. To address this issue, we propose decomposing client LoRA updates via Robust Principal Component Analysis (Robust-PCA) into a common low-rank component and client-specific sparse components. Our proposed algorithm FedRPCA aggregates the low-rank components through averaging, consolidating common knowledge, and applies scaled averaging to the sparse components to amplify client-specific knowledge. We evaluate our approach across a variety of vision and language tasks and demonstrate that it achieves higher final accuracy and faster convergence compared to competing baselines.