Rethinking LoRA for Data Heterogeneous Federated Learning: Subspace and State Alignment
This addresses the problem of efficient and robust federated fine-tuning for clients with heterogeneous data, representing an incremental improvement over existing methods.
The paper tackled the performance gap of Low-Rank Adaptation (LoRA) in federated learning under non-IID data by identifying update-space and optimizer-state mismatches, and proposed FedGaLore, which improved robustness and accuracy over state-of-the-art baselines across NLU, vision, and NLG benchmarks.
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.