Preventing Rank Collapse in Federated Low-Rank Adaptation with Client Heterogeneity

TL;DR

This paper identifies a phenomenon called rank collapse in federated low-rank adaptation with heterogeneous client ranks, analyzes its root cause, and proposes a new aggregation method, raFLoRA, to prevent it and improve performance.

Contribution

The paper reveals the cause of rank collapse in FedLoRA and introduces raFLoRA, a novel rank-partitioned aggregation method that enhances robustness and performance.

Findings

raFLoRA prevents rank collapse across diverse tasks

It improves model performance compared to baseline FedLoRA methods

It enhances robustness in heterogeneous federated learning settings

Abstract

Federated low-rank adaptation (FedLoRA) has facilitated communication-efficient and privacy-preserving fine-tuning of foundation models for downstream tasks. In practical federated learning scenarios, client heterogeneity in system resources and data distributions motivates the use of heterogeneous LoRA ranks across clients. However, we identify a previously overlooked phenomenon in heterogeneous FedLoRA with SVD-based allocation, termed rank collapse, where the energy of the global update becomes concentrated in the minimum shared rank, resulting in suboptimal performance and high sensitivity to rank configurations. Through theoretical analysis, we reveal the root cause of rank collapse: a mismatch between rank-agnostic aggregation weights and rank-dependent client contributions, which systematically suppresses higher-rank updates at a geometric rate over rounds. Motivated by this insight, we propose raFLoRA, a rank-partitioned aggregation method that decomposes local updates into rank partitions and then aggregates each partition weighted by its effective client contributions. Extensive experiments across vision, language, and reasoning tasks show that raFLoRA prevents rank collapse, improves model performance, and enhances robustness across diverse heterogeneous configurations compared with strong FedLoRA baselines.