Preventing Rank Collapse in Federated Low-Rank Adaptation with Client Heterogeneity
This addresses a specific bottleneck in federated learning for fine-tuning foundation models, offering an incremental improvement for scenarios with heterogeneous client resources and data.
The paper tackled the problem of rank collapse in federated low-rank adaptation (FedLoRA) due to client heterogeneity, where performance degrades as updates concentrate on the minimum shared rank. The result was a proposed method, raFLoRA, which prevented rank collapse and improved model performance across classification and reasoning tasks while maintaining communication efficiency.
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 heterogeneous LoRA ranks across clients. We identify a previously overlooked phenomenon in heterogeneous FedLoRA, termed rank collapse, where the energy of the global update concentrates on 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 classification and reasoning tasks show that raFLoRA prevents rank collapse, improves model performance, and preserves communication efficiency compared to state-of-the-art FedLoRA baselines.