Rethinking Factor Sharing in Federated LoRA: A Rank-Aware Adaptive Approach

This paper investigates which of the two low-rank adaptation (LoRA) factors, A or B, should be shared across clients in federated learning, and which should remain client-specific. The authors compare two strategies: Share-A/Local-B and Share-B/Local-A. Through a least-squares surrogate analysis, they reveal that Share-A/Local-B imposes a common rank-r input-side space across clients, while Share-B/Local-A imposes a common rank-r output-side space. These strategies incur different projection residuals, and the preferred strategy is the one with the smaller aggregate residual across clients.

Based on this insight, the authors propose Federated Adaptive Factor Sharing Low-Rank Adaptation (FedAS-LoRA), which selects the sharing side before training. To enable this selection, they design a Rank-Aware Shared-Subspace Sufficiency (RSS) metric that uses sequence-level representations from a frozen LLM backbone to assess whether a global rank-r input subspace is sufficient for local data distributions. The RSS metric compares the representation energy retained by a global rank-r subspace with that retained by client-specific rank-r subspaces, and a threshold is calibrated using random-split and bootstrap procedures. If the RSS deficit exceeds the threshold, Share-B/Local-A is selected; otherwise, Share-A/Local-B is selected.

The authors prove convergence of FedAS-LoRA under arbitrary client participation for both sharing strategies, showing that the iterates converge to a stationary neighborhood. Experiments are conducted on natural language understanding tasks (GLUE benchmark with RoBERTa-large) and natural language generation (GSM8K with LLaMA3-8B). Results show that FedAS-LoRA consistently outperforms fixed sharing strategies and other federated LoRA baselines across different tasks, data distributions, LoRA ranks, and participation settings. The RSS-based selection is validated to align with the empirically better sharing strategy in all 33 cases with a strict accuracy difference out of 38 evaluated settings. The paper concludes that no fixed sharing policy is universally optimal, and the preferred sharing side depends on the data distribution, adapter rank, and participation pattern.

Improvements to AI Systems:

1. Adaptive Parameter-Efficient Fine-Tuning for Federated Learning:

• The improved system can dynamically choose whether to share the input-side or output-side low-rank adapter factor across clients based on a pre-training metric (RSS). This eliminates the need for manual tuning of sharing strategies, which is currently a major bottleneck in federated LoRA.

• The system can automatically detect when client-specific output spaces are more beneficial (e.g., heterogeneous label distributions) versus when shared input spaces suffice (e.g., homogeneous data), leading to faster convergence and higher final accuracy without exhaustive experimentation.

2. Data-Distribution-Aware Model Personalization:

• The improved system can use the RSS metric to quantify how well a global rank-r subspace captures local data representations before training begins. This enables proactive selection of the optimal sharing side, reducing communication overhead and improving robustness to non-IID data.

• For example, in a federated medical imaging scenario where each hospital has different imaging protocols, the system would automatically choose Share-B/Local-A to preserve client-specific output mappings, avoiding catastrophic forgetting of rare local patterns.

3. Rank-Adaptive Federated Learning:

• The system can now calibrate the threshold for RSS deficits using bootstrap procedures, allowing it to adaptively switch between sharing strategies as the rank r changes. This is crucial for resource-constrained edge devices where the optimal rank varies with memory and compute budgets.

• The improved system can also provide a confidence score for the chosen strategy, enabling safe deployment in production where a wrong choice could degrade performance.

4. Convergence-Guaranteed Federated Optimization:

• The improved system inherits the theoretical convergence guarantee under arbitrary client participation, meaning it can handle realistic scenarios with clients dropping in and out. This makes it suitable for large-scale cross-device federated learning (e.g., mobile keyboards, smart home devices) where participation is unpredictable.

• The system can now reliably train large language models (e.g., LLaMA3-8B) on decentralized data with minimal communication cost, achieving accuracy comparable to centralized fine-tuning.

5. Generalizable Fine-Tuning for Heterogeneous Tasks:

• The improved system can be applied to both natural language understanding (e.g., GLUE) and generation (e.g., GSM8K) without modification, thanks to the sequence-level representation-based RSS metric that is model-agnostic. This enables a single federated framework to serve diverse applications—from sentiment analysis to math reasoning—while automatically adapting to each task’s data geometry.

6. Zero-Shot Strategy Selection for New Federated Clients:

• When a new client joins, the system can use the RSS metric on that client’s local data (using the frozen backbone) to decide whether to keep the current global sharing strategy or switch. This allows for online adaptation without retraining, improving scalability and reducing downtime in dynamic federated networks.

What the Improved AI System Can Do:

• Automatically determine the optimal LoRA sharing strategy for any federated learning task, dataset, rank, and participation pattern, without manual tuning.

• Achieve higher accuracy and faster convergence than fixed-sharing baselines, especially under non-IID data distributions.

• Operate with theoretical convergence guarantees even with unreliable client participation.

• Deploy on resource-constrained devices by using low-rank adapters and minimal communication overhead.

• Provide a principled, data-driven method to balance personalization (client-specific factors) and generalization (shared factors) in federated fine-tuning of large language models.

Low-rank adaptation (LoRA) represents large language model (LLM) updates with two compact matrix factors, i.e., A and B, providing an efficient way to fine-tune large models in federated learning paradigm. Inspired by the asymmetric roles of the LoRA factors, we study whether A should be shared across clients while B remains client-specific (Share-A/Local-B), or whether B should instead be shared while A remains client-specific (Share-B/Local-A). With a least-squares surrogate, we reveal that Share-A/Local-B requires the client-specific LoRA update matrices to use a common rank- r input-side space, whereas Share-B/Local-A requires a common rank- r output-side space. The two strategies therefore incur different projection residuals, indicating that the preferred strategy is the one with the smaller aggregate residual across clients. With this insight, we propose Federated Adaptive Factor Sharing Low-Rank Adaptation (FedAS-LoRA), which selects the sharing side before training to enhance fine-tuning performance. To enable adaptive factor selection before training, we design a Rank-Aware Shared-Subspace Sufficiency (RSS) metric, which effectively assesses whether a shared rank- r input subspace is sufficient for the local data distributions using representations extracted from a frozen LLM backbone. Experiments across different tasks, data distributions, LoRA ranks, and participation settings confirm the effectiveness of RSS and the superior performance of FedAS-LoRA.