SeFoRA: Sketch-Aggregated Federated Low-Rank Adaptation with Heterogeneous Client Ranks

SeFoRA: Sketch-Aggregated Federated Low-Rank Adaptation with Heterogeneous Client Ranks

Abstract

The paper considers federated parameter-efficient fine-tuning of large neural networks with low-rank adaptation (LoRA). Combining LoRA with federated PEFT introduces challenges absent from either setting alone: clients may use different LoRA ranks, making their factor matrices dimension-incompatible, and factor-wise averaging suffers from a bilinear mismatch. The authors propose SeFoRA, a sketch-aggregated federated LoRA algorithm in which each client transmits a linear sketch of its local updates, enabling direct aggregation at the federator. As a result, SeFoRA alleviates the bilinear mismatch and allows for aggregation in a small subspace of the full model. They introduce a rank-homogeneous version called SeFoRA-Ho which allows for direct adapter aggregation in this setting. They prove convergence to a neighborhood of the first-order stationary point at rate O(1/T) for the rank-homogeneous setting. Numerical experiments on fine-tuning RoBERTa-Large on GLUE datasets show how their algorithms outperform the state-of-the-art.

Introduction and Motivation

Foundation models such as GPT-4, BERT, PaLM 2, Llama 2, Claude 3, and ViT have remarkable learning capabilities across diverse tasks. Adapting such models to downstream tasks requires vast datasets and tremendous computational resources. Parameter-efficient fine-tuning (PEFT) methods adapt a pretrained model by updating only a small, structured subset of parameters. Federated learning (FL) is a distributed learning paradigm where a central server coordinates training on heterogeneous data owned by participating clients, keeping data local and sharing only model updates.

Combining PEFT with FL spans adapter-based methods, prompt tuning, selective parameter optimization, and low-rank gradient subspace optimization. The paper focuses on federated PEFT methods using LoRA, which has emerged as one of the most popular PEFT techniques.

Combining LoRA with FL introduces challenges not present in either setting alone. In LoRA, for each layer, each client computes an update ∆Wi = BiAi, where the model update ∆Wi ∈ R(d×s) is represented as a product of two factor matrices Bi ∈ R(d×r) and Ai ∈ R(r×s) for a chosen r ≪ min(d,s). The server aims to obtain ∆W = Σi pi ∆Wi = Σi pi BiAi. The paper identifies two main challenges:

Challenge 1: Rank-heterogeneity – Clients use different adapter sizes ri to match local compute budgets or data complexity, resulting in matrices with incompatible dimensions in aggregation. Allowing rank heterogeneity is crucial since small adapter ranks can be sufficient for certain tasks but may not always work, especially when the downstream task differs significantly from pretraining.

Challenge 2: Bilinear mismatch – Zero-padding and averaging adapters first has lower computational cost but does not recover the correct average of updates: BA ≠ Σi pi BiAi. This problem is exacerbated when considering differential privacy guarantees and the non-linearity of aggregation does not allow for secure aggregation.

Key Contributions

1. SeFoRA aggregates fixed-dimensional sketches, irrespective of clients' local ranks, and can be seamlessly combined with adaptive rank-truncation at clients.

2. A rank-homogeneous version called SeFoRA-Ho is provided with convergence guarantee on the original LoRA objective at rate O(1/T + c), where c is a sketching-dependent error floor bounded using spectral tail energy of aggregated client updates.

3. Numerical experiments fine-tuning RoBERTa-Large on GLUE datasets in rank-heterogeneous and rank-homogeneous settings show SeFoRA and SeFoRA-Ho outperform existing methods.

4. Mitigating the bilinear mismatch and allowing for linear aggregation open the door for secure aggregation and more efficient differential privacy mechanisms.

Algorithm Description

SeFoRA resolves rank-heterogeneity and bilinear mismatch by operating in a sketched-space of the full model using linear matrix sketching. Each client sketches its local update ∆Wi into two compact, fixed-size matrices using a sketching method. Since the sketching operator is linear, aggregating the sketches at the server is equivalent to computing the sketch of the aggregated update exactly, overcoming the bilinear mismatch with low computational overhead.

The procedure works as follows:

• Each client sets its own local rank profile ri = (ri,1,..., ri,L)

• Clients can periodically update ranks using Algorithm 4 (Update Local Rank)

• After local fine-tuning, clients compute sketch matrices Yi = ∆Wiomega and Zi = Ψ∆Wi

• The federator aggregates sketches linearly: Y = Σi piYi and Z = Σi piZi

• The federator uses Algorithm 1 (Unsketching) to recover low-rank factors B and A

SeFoRA-Ho is the rank-homogeneous version where all clients use the same rank, obtained by modifying Step 7 in Algorithm 2 and removing Step 9.

Theoretical Analysis

Under Assumptions 1-3 (per-sample Lipschitz smoothness, unbiased stochastic gradients with bounded moments, bounded LoRA factors), Theorem 1 shows SeFoRA-Ho converges. The convergence bound is:

(1/T) Σt E[∇a L lora(B(t), A(t))2 F + ∇ B L lora(B(t), A(t))2 F] ≤ 2∆0/(ηET) + (2/(ηE))S1 + ηE·2(χC3 + 2C4) + η2E2·χ2(C1 + C2) + 4η3E3·µC32

where C1 = χ2Ma2 + 2µ2(M B4 + Ma4), C2 = χ2M B2 + 2µ2Ma2(M B4 + Ma4), C3 = Ma2M B2κ4, C4 = µχ2(M B2 + Ma2), S1 is the sketching floor, and ∆0 is the initial optimality gap.

Corollary 1 shows that with bounded spectral tail (Assumption 4), SeFoRA-Ho converges at rate O(1/T + c) where c is a sketching-dependent error floor.

Experimental Results

The paper evaluates SeFoRA-Ho in rank-homogeneous settings and SeFoRA in rank-heterogeneous settings using RoBERTa-Large on four GLUE datasets (SST-2, QNLI, QQP, MNLI). They use 40 clients, Dirichlet distribution with α = 0.5 for data heterogeneity, local steps E ∈ 1, 5, Adam optimizer with learning rate 9e-4, and maximum adapter rank r = 16 with sketch dimension k = 18.

Rank-homogeneous results: With E=5, SeFoRA-Ho achieves 95.7% on SST-2, 93.7% on QNLI, 87.9% on QQP, and 89.6% on MNLI, outperforming FedIT and FFA-LoRA.

Rank-heterogeneous results: SeFoRA achieves performance comparable to or slightly better than FlexLoRA without the need to reconstruct all clients' updates, which incurs O(d2s) computational cost at the server. FLoRA converges considerably more slowly, particularly on QNLI.

Sensitivity analysis: For fixed rank r, increasing sketch dimension from k = r+2 to 2r produces little or no improvement, indicating k = r+2 is sufficient. Small ranks are sufficient for SST-2 and QNLI, while more challenging tasks like QQP and MNLI benefit from larger adapter ranks.

Rank adaptation: The average local rank decreases substantially across communication rounds, showing adapter updates demand reduced rank as federated fine-tuning progresses. Rank adaptation reflects differences in task complexity.

Related Work

The paper provides extensive review of federated LoRA methods, including FedIT, FLoRA, FlexLoRA, HetLoRA, FSLoRA, FFA-LoRA, and others, discussing their approaches to handling bilinear mismatch and rank heterogeneity. The key comparison table shows SeFoRA uniquely mitigates bilinear mismatch without operating in full model space, with server compute cost O(Nr(d+s)) and convergence rate O(1/T + c).

Based on this paper, here are the specific improvements I can make to AI systems:

• What I can do: Enable federated fine-tuning where clients (edge devices, organizations) use different adapter ranks based on their local compute budgets, memory constraints, or data complexity—without requiring rank homogenization or full-model reconstruction at the server.

• Improved system capability: A mobile app ecosystem can collaboratively fine-tune a shared language model where phones with 2GB RAM use rank-4 adapters, tablets with 4GB use rank-8, and servers use rank-16—all contributing to a single improved model without any client being forced into an incompatible configuration.

• What I can do: Replace non-linear adapter averaging (which breaks secure aggregation and complicates differential privacy) with linear sketch aggregation. This makes the aggregation step compatible with cryptographic secure aggregation protocols and standard DP noise mechanisms.

• Improved system capability: A healthcare consortium can fine-tune a diagnostic model across hospitals using LoRA adapters, with formal differential privacy guarantees (e.g., ε-DP) applied directly to the linear sketches—something impossible with existing bilinear-mismatch-prone methods. The system can now provably protect patient data while still achieving state-of-the-art accuracy.

• What I can do: Dynamically adjust each client's LoRA rank during federated training based on the evolving task difficulty and data complexity, using the rank-update algorithm (Algorithm 4). The system detects when adapter updates become low-rank and shrinks ranks to save computation, or expands them when needed.

• Improved system capability: A federated recommendation system can start with rank-16 adapters for all users, then automatically reduce to rank-4 for users with simple, repetitive behavior patterns (saving 75% of their local compute) while keeping rank-16 for power users with diverse interests—all without server-side intervention or manual tuning.

• What I can do: Aggregate client updates in a small sketched subspace (dimension k ≈ r+2) rather than reconstructing full model updates, reducing server computation from O(d2s) to O(Nr(d+s)) where N is the number of clients.

• Improved system capability: A central server coordinating 10,000 edge devices can now aggregate updates in milliseconds instead of minutes, enabling real-time federated fine-tuning for latency-sensitive applications like voice assistants or autonomous vehicle fleets that require frequent model refreshes.

• What I can do: Provide provable convergence to a neighborhood of the stationary point at rate O(1/T + c), where the error floor c is explicitly bounded by the spectral tail energy of the aggregated updates. This gives system designers a theoretical guarantee on how close the final model will be to optimality.

• Improved system capability: A financial fraud detection system can now certify that after T communication rounds, the fine-tuned model is within a known, computable distance from the optimal solution—enabling compliance with regulatory requirements that demand verifiable model quality guarantees.

• What I can do: Use the finding that sketch dimension k = r+2 is sufficient (with negligible gains beyond 2r) to minimize communication overhead. Each client transmits only 2k(d+s) numbers instead of the full d×s update matrix.

• Improved system capability: A satellite-based federated learning system with severely bandwidth-constrained links can now fine-tune a vision model for crop monitoring by transmitting only 0.1% of the model update size, while maintaining 95%+ of the performance of full-update methods—enabling practical deployment in remote or low-connectivity regions.

• What I can do: Leverage the observation that different tasks require different adapter ranks (SST-2/QNLI work well with small ranks; QQP/MNLI need larger ones) to automatically assign ranks per client based on their local task's inherent complexity.

• Improved system capability: A federated NLP platform serving multiple downstream tasks (sentiment analysis, question answering, paraphrase detection, natural language inference) can automatically allocate more computational resources to clients working on harder tasks, while simpler tasks run efficiently on minimal hardware—optimizing the overall system's energy consumption and throughput.