Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Hongyi Peng, Han Yu, Xiaoxiao Li, Qiang Yang

TL;DR

Non-IID data causes federated LoRA to underperform full fine-tuning due to update-space and optimizer-state drift. FedGaLore merges GaLore-style gradient-subspace client optimization with AJIVE-based synchronization of projected second-moment states to address both failure modes, delivering robust performance close to FFT with PEFT efficiency. The approach yields consistent improvements across NLP, vision, and LLM benchmarks and remains communication-light by transmitting only rank-$r$ projections and synchronized statistics. This work enhances the reliability of PEFT in heterogeneous FL, enabling scalable fine-tuning of foundation models without full-parameter updates.

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.

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

TL;DR

Non-IID data causes federated LoRA to underperform full fine-tuning due to update-space and optimizer-state drift. FedGaLore merges GaLore-style gradient-subspace client optimization with AJIVE-based synchronization of projected second-moment states to address both failure modes, delivering robust performance close to FFT with PEFT efficiency. The approach yields consistent improvements across NLP, vision, and LLM benchmarks and remains communication-light by transmitting only rank- projections and synchronized statistics. This work enhances the reliability of PEFT in heterogeneous FL, enabling scalable fine-tuning of foundation models without full-parameter updates.

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.
Paper Structure (52 sections, 5 theorems, 85 equations, 6 figures, 11 tables, 5 algorithms)

This paper contains 52 sections, 5 theorems, 85 equations, 6 figures, 11 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Preliminaries
  4. Notation.
  5. Problem Setup
  6. Assumptions
  7. A W.H.P. Lens for Federated Learning
  8. Why W.H.P. robustness?
  9. Update-space Mismatch
  10. Federated LoRA differs from FFT.
  11. Optimizer-state Mismatch
  12. The Proposed FedGaLore Method
  13. Client-side optimization (GaLore on target modules).
  14. Server-side: FedAvg for parameters, AJIVE for second-moment synchronization.
  15. Computation and communication.
  16. ...and 37 more sections

Key Result

Lemma 4.1

Let $Q$ be convex. If $\theta_i\in Q$ for all $i\in[M]$ and $\bar{\theta}=\sum_{i=1}^M \tilde{p}_i \theta_i$ with $\tilde{p}_i\ge 0$ and $\sum_i \tilde{p}_i=1$, then $\bar{\theta}\in Q$.

Figures (6)

  • Figure 1: Left: Under increasing data heterogeneity (smaller Dirichlet $\alpha$), representative federated LoRA baselines degrade sharply, while FedGaLore remains stable and approaches full fine-tuning performance. Right: With local adaptive optimizers, client training loss can decrease while global validation performance stagnates or degrades, indicating optimizer-state mismatch.
  • Figure 2: Trajectory robustness and failure modes. The stable region $Q$ (blue tube) surrounds a centralized reference solution $\theta^\star$ (dashed). Left (round $k$): Stable operation; all client trajectories remain in $Q$. Middle (round $k{+}1$):Local escape ($\mathcal{E}_{\mathrm{traj}}^{c}$, red), where client drift causes a trajectory to exit $Q$. Right (round $k{+}2$):Aggregation instability ($\mathcal{E}_{\mathrm{agg}}^{c}$), where the aggregated model leaves $Q$ even though all client endpoints lie in $Q$.
  • Figure 3: Geometric robustness and aggregation stability.(a--b) Synthetic landscape: GaLore-style updates reach flat basins more often than fixed-subspace LoRA (60% vs. 20%; Appendix \ref{['app:synthetic_exp']}). (c) Non-IID ViT-MNIST: linear interpolation between client models shows low loss barriers for FFT and GaLore-style updates, while LoRA exhibits a higher barrier, indicating reduced aggregation robustness.
  • Figure 4: Local convergence under different projector schedules. FedGaLore (SVD$\rightarrow$random) achieves the best time-to-loss compared to always-SVD and always-random schedules.
  • Figure 5: AJIVE recovers global geometry under structured drift. We compare the reconstruction error of the global second moment $V^\star$ ($y$-axis) vs. number of clients ($x$-axis). Naive Averaging (Red) fails to converge to the true signal because the squared drift terms $\mathbb{E}[\bm{L}_k^2]$ introduce a persistent bias. SVD on Average (Solid lines) helps slightly but operates on already-corrupted data. AJIVE (Markers) successfully isolates the shared signal structure. Notably, AJIVE Rank-15 (Green circles) achieves the lowest error, correctly capturing the rank expansion induced by the element-wise square operation.
  • ...and 1 more figures

Theorems & Definitions (12)

  • Definition 3.1: Local Training Operator
  • Definition 3.2: Server Aggregation Operator
  • Definition 3.3: State Synchronization Protocol
  • Lemma 4.1: Aggregation preserves convex regions
  • Theorem 4.2: Weyl's tube formula (informal) gray2003tubes
  • Definition 4.3: Biased state initialization
  • Theorem 4.4: W.H.P. local-containment radius
  • Definition 2.1: Conditional sub-Gaussian vector
  • Lemma 2.2: W.H.P. noise envelope
  • proof
  • ...and 2 more