Rethinking LoRA for Data Heterogeneous Federated Learning: Subspace and State Alignment

TL;DR

This paper identifies key mismatches limiting LoRA's performance in non-IID federated learning and proposes FedGaLore, a novel method combining gradient subspace optimization and drift-robust synchronization, to enhance robustness and accuracy.

Contribution

It introduces FedGaLore, a new federated learning approach that addresses update and optimizer-state mismatches in LoRA under non-IID data distributions.

Findings

FedGaLore outperforms existing federated LoRA methods in non-IID settings.

Improves robustness and accuracy across NLP, vision, and NLG benchmarks.

Addresses update-space and optimizer-state mismatches effectively.

Abstract

Low-Rank Adaptation (LoRA) is widely used for federated fine-tuning. Yet under non-IID settings, it can substantially underperform full-parameter fine-tuning. Through with-high-probability robustness analysis, we uncover that this gap can be attributed to two coupled mismatches: (i) update-space mismatch, where clients optimize in a low-rank subspace but aggregation occurs in the full space; and (ii) optimizer-state mismatch, where unsynchronized adaptive states amplify drift across rounds. We propose FedGaLore, which combines client-side GaLore-style gradient-subspace optimization with server-side drift-robust synchronization of projected second-moment states via spectral shared-signal extraction, to address this challenge. Across NLU, vision, and NLG benchmarks, FedGaLore improves robustness and accuracy over state-of-the-art federated LoRA baselines in non-IID settings.