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

2026-02-02

AI Summary

In non-IID federated learning, Low-Rank Adaptation (LoRA) suffers significant performance degradation due to mismatches between client update subspaces and the server aggregation space, as well as desynchronization of optimizer states. This work proposes FedGaLore, a novel framework that, for the first time, identifies and decouples these two sources of mismatch. It achieves gradient subspace alignment through spectral analysis–driven shared signal extraction and introduces a projected second-moment state synchronization mechanism. By integrating GaLore-style gradient optimization with robust state synchronization, FedGaLore substantially improves the accuracy and robustness of LoRA under data heterogeneity across multiple benchmarks in natural language understanding, vision, and natural language generation, outperforming existing state-of-the-art methods.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Natural Language Processing: Language Grounding & Multi-modal NLP
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Federated Learning
• LoRA
• non-IID
• Update-space Mismatch
• Optimizer-state Mismatch

Innovation

• LoRA
• Federated Learning
• non-IID
• Subspace Optimization
• Optimizer State Synchronization