🤖 machine learning

Subspace-Constrained Federated Learning with Low-Rank Adaptation

This paper proposes a subspace-regularized federated learning objective to mitigate geometric misalignment in heterogeneous client data during LoRA fine-tuning, demonstrating through extensive experiments on RoBERTa-large and SmolLM-360M that this approach significantly improves convergence and accuracy by enforcing near-perfect basis overlap across clients.

Original authors: Neranjan Senarath, Rohit Muralitharan, Sadia Asif

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Neranjan Senarath, Rohit Muralitharan, Sadia Asif

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of 10 friends trying to learn a new dance routine together, but they are all in different rooms (this is Federated Learning). They can't share their private video recordings of themselves practicing (privacy), and they can only send short text messages describing their moves to a central coach (communication limits).

To save time and bandwidth, they don't send their whole dance routine. Instead, they only send a few key "moves" or "tweaks" they've invented to make the dance better. This is similar to LoRA (Low-Rank Adaptation), a technique that lets people learn big tasks by only adjusting a tiny, efficient part of the model.

The Problem: Dancing in Different Directions

The paper identifies a hidden problem: because each friend is practicing to a different song or style (heterogeneous data), they end up inventing moves that point in completely different directions.

Friend A invents a move that spins left.
Friend B invents a move that spins right.
Friend C invents a move that jumps forward.

When the coach collects all these moves to create one "Global Dance," they cancel each other out. The left spin cancels the right spin; the jump gets lost in the confusion. The paper calls this "Silent Cancellation." The coach sees a lot of effort from everyone, but the final result is weak because the moves are fighting against each other.

The Solution: The "Shared Reference" Rule

The authors propose a new rule called Subspace-Regularized LoRA.

Imagine the coach gives everyone a shared reference video of the "ideal" move before they start practicing. The rule is: "You can invent your own new moves, but they must stay very close to the direction of this shared reference video."

This is the Subspace Constraint. It doesn't stop them from learning; it just gently nudges them to keep their "dance floor" aligned with everyone else's.

Instead of Friend A spinning wildly left and Friend B spinning wildly right, they both adjust their spins to be closer to the coach's reference direction.
When the coach combines their moves, they reinforce each other instead of canceling out.

What They Tested

The researchers tested this idea on two different "dancers" (AI models):

RoBERTa-large: A very large, complex model.
SmolLM-360M: A smaller, more compact model.

They simulated 10 friends (clients) with different data and ran the experiment 24 times to be sure.

The Results

1. The "Alignment" Score (Did they dance together?)
The team measured how well the friends' moves aligned.

Old methods: The friends' moves were somewhat aligned (about 96% to 99% overlap).
New method: The friends' moves were almost perfectly aligned (99.99% overlap). They were all dancing on the exact same "floor."

2. The Performance Score (Did the dance look good?)

On the Large Model (RoBERTa): The new method was a huge success. Because the moves were perfectly aligned, the final dance was much better than before. The accuracy jumped significantly, and the "mistakes" (loss) went down.
On the Small Model (SmolLM): Here, the result was surprising. Even though the new method achieved that perfect 99.99% alignment, the final dance score didn't improve much compared to the old methods. The small model seemed to do just fine even with a little bit of misalignment.

The Big Takeaway

The paper proves that geometric alignment (making sure everyone is moving in the same general direction) is real and can be achieved with their new rule.

However, the paper also warns that perfect alignment doesn't always guarantee a better final score.

For the big, complex model, alignment was the key to unlocking better performance.
For the smaller model, alignment happened perfectly, but it didn't translate into a massive score boost.

In short: The authors built a system that forces AI models to "agree" on the direction of their learning updates. This stops them from fighting each other. It works brilliantly for big models, but for smaller models, just agreeing on direction isn't the only thing that matters for getting a high score.

Technical Summary: Subspace-Constrained Federated Learning with Low-Rank Adaptation

1. Problem Statement

Federated Learning (FL) enables collaborative training across decentralized clients without sharing raw data, yet it faces significant challenges when client data distributions are heterogeneous. In standard Federated Averaging (FedAvg), statistical heterogeneity leads to "client drift," where local model updates diverge, causing instability and slow convergence.

With the rise of large foundation models, Low-Rank Adaptation (LoRA) has become a popular parameter-efficient fine-tuning technique for FL. LoRA constrains weight updates ( $\Delta W$ ) to a low-rank decomposition ($BA$), reducing communication and computational costs. However, this paper hypothesizes that under heterogeneous data distributions, different clients learn low-rank updates that span distinct subspaces. When these geometrically misaligned updates are naively aggregated at the server, they may partially cancel each other out—a phenomenon termed "silent cancellation." This structural cancellation reduces the effective magnitude of the global update, hindering convergence even when individual clients make substantial local progress.

2. Methodology

The authors propose Subspace-Regularized LoRA (Subspace-Reg), a framework designed to align local client updates with a shared global reference subspace.

2.1 Objective Function

Each client $k$ optimizes a subspace-regularized objective that includes a proximal alignment penalty. The objective is:
$L_k^{SR}(w_k) = \min_{A_k, B_k} \left\{ L_k(w_k) + \frac{\mu}{2} \|B_k A_k - \Delta W_g^{(t)}\|_F^2 + \frac{\lambda}{2} (\|A_k\|_F^2 + \|B_k\|_F^2) \right\}$
Where:

$L_k(w_k)$ is the local loss.
$\Delta W_g^{(t)}$ is the current global low-rank update (server reference).
$\mu$ controls the strength of subspace alignment.
$\lambda$ regularizes the low-rank factors.
The Frobenius norm term constrains local updates to remain close to the server's reference update, encouraging constructive aggregation.

2.2 Algorithm

The training follows a standard FL structure with modifications:

Initialization: The server initializes a global low-rank basis $B_g$ with orthonormal columns.
Server Step: The server samples clients, computes a global reference update via SVD of the weighted sum of previous client updates, and broadcasts the model and reference basis.
Client Step: Each client initializes their local factors to the global reference and performs $E$ $E$ local epochs. The local gradient descent minimizes the subspace-regularized loss, which includes:
- The primary task loss.
- A term penalizing deviation of the product $B_k A_k$ from the global reference $B_g A_g$ .
- An optional term penalizing deviation of the basis $B_k$ from $B_g$ .
Aggregation: Clients send their updates ( $\Delta W_k = B_k A_k$ ) to the server, which aggregates them and updates the global model.

3. Theoretical Analysis

The paper provides a theoretical analysis to understand how subspace alignment influences convergence under heterogeneity.

Geometric Cancellation Bound: The authors derive that the squared norm of the aggregated update depends on the cross-terms between client updates. Under heterogeneity, if updates are misaligned, these cross-terms become negative, reducing the effective update magnitude (silent cancellation).
Convergence Guarantees: Under standard assumptions (smoothness, bounded heterogeneity, bounded variance, and bounded low-rank updates), the paper proves a convergence bound for the average squared gradient norm.
- The bound includes a term $C_1 \sigma^2 / \mu^2$ representing the misalignment penalty.
- Unlike standard LoRA-FedAvg, which suffers from a fixed cancellation term, Subspace-Reg allows this penalty to be made arbitrarily small by increasing the alignment parameter $\mu$ .
- The authors derive conditions for $\mu$ and the number of communication rounds $T$ to achieve a target gradient accuracy $\epsilon$ .

4. Experimental Results

The method was evaluated on two pretrained models (RoBERTa-large and SmolLM-360M) using the HellaSwag dataset in a non-IID setting (10 clients, Dirichlet $\beta=0.5$ ). Experiments were run across 3 random seeds (42, 43, 44), totaling 24 runs.

4.1 Performance on RoBERTa-large

Subspace-Reg significantly outperformed all baselines (FedAvg, SVD redistribution, and FedSVD):

Accuracy: Achieved the highest mean best accuracy (0.454 ± 0.023) and mean final accuracy (0.429 ± 0.011), outperforming the best baseline (FedSVD) by a margin of ~0.118 in best accuracy.
Loss: Achieved the lowest final loss (1.363).
Alignment: Achieved near-perfect Basis Overlap (≈0.9999), compared to 0.958–0.991 for baselines.

4.2 Performance on SmolLM-360M

The results on SmolLM-360M revealed that accuracy gains are model-dependent:

Accuracy: FedAvg achieved the highest accuracy (0.581), while Subspace-Reg ranked second (0.564).
Alignment: Despite the lack of accuracy improvement, Subspace-Reg maintained near-perfect Basis Overlap (0.99989), consistent with the RoBERTa results.
Conclusion: The geometric alignment was successful, but it did not translate to downstream accuracy gains for this specific model.

4.3 Key Diagnostic Findings

Basis Overlap vs. Cosine Similarity: While SVD-based methods achieved higher cosine-based metrics (Client-Global Cosine, Pairwise Cosine), they did not achieve the same accuracy improvements as Subspace-Reg on RoBERTa. The authors identify Basis Overlap (direct subspace alignment) as the critical diagnostic, rather than cosine similarity.
Robustness: The near-perfect basis overlap was consistent across all seeds and both models, validating the geometric alignment hypothesis even when downstream performance varied.

5. Significance and Claims

The paper claims the following contributions and significance:

Identification of Silent Cancellation: It highlights that subspace misalignment in LoRA-based FL leads to destructive aggregation, a phenomenon previously underexplored.
Subspace-Regularized Objective: It proposes a proximal alignment penalty that explicitly encourages local updates to remain close to a global reference subspace, extending the concept of FedProx from parameter magnitude to geometric structure.
Theoretical Improvement: It provides a convergence bound showing that the misalignment penalty is tunable, offering a strictly better guarantee than standard LoRA-FedAvg under heterogeneity.
Empirical Validation: It demonstrates that the proposed method achieves near-perfect basis overlap (≈0.9999) consistently. On RoBERTa-large, this translates to substantial accuracy gains. On SmolLM-360M, it confirms that geometric alignment can be achieved independently of downstream accuracy, suggesting that the benefits of alignment are model-dependent.
Diagnostic Insight: The work establishes that basis overlap is a more relevant indicator of aggregation quality than cosine similarity in federated LoRA settings.

The authors conclude that while subspace alignment is a robust geometric mechanism, its impact on final task performance depends on the specific model architecture and capacity, warranting further investigation into why these benefits vary across models.