Align Before Aggregation: Basis-Consistent Federated LoRA under Heterogeneous Ranks

AI summary

Overview

Research area: Federated learning for large language models, specifically parameter-efficient fine-tuning with low-rank adaptation (LoRA) when different clients use different LoRA ranks.

Technical level: Advanced. The core argument is intuitive, but the method and theory rely on matrix factorization, principal angles, orthogonal Procrustes alignment, and SVD-based refactorization.

Scope: The paper diagnoses a "basis inconsistency" problem in server-side aggregation of heterogeneous-rank LoRA updates and proposes FedAbA, a server-side framework that aligns every client to a shared global reference basis before aggregation, with supporting theory and experiments on NLU, NLG-style instruction tuning, and natural-instructions benchmarks.

What This Paper Is About

Clients in federated learning have different compute and bandwidth budgets, so they naturally train LoRA adapters of different ranks. Servers typically compose each client's two small LoRA matrices into a full matrix update, average those matrices, refactorize the average with SVD, and send back the first r_i rows/columns to each client. Because a low-rank factorization is not unique — B_iA_i = (B_iR)(R^{-1}A_i) for any invertible R — each client's factors can live in a private "basis," and the rank-r_i prefix a low-rank client receives may mean something different from round to round. The goal is to make aggregation consistent against one explicit global reference basis, and to weight clients by both data size and how well their local basis matches that reference.

Key Contributions

1. Problem diagnosis. The authors identify that the real obstacle in heterogeneous-rank federated LoRA is not a dimension mismatch — every ΔW_i lives in the same ℝ^{m×n} — but the absence of explicit basis reconciliation against a global reference basis in the compose–aggregate–factorize–dispatch loop. They also prove (Proposition 3.1) that merely rotating local factors is a no-op, since sB̄_iT(Q_i^tT)^⊤ = sB̄_i(Q_i^t)^⊤.

2. The FedAbA framework. A server-side pipeline that (i) extracts each client's local basis by thin QR of (A_i^t)^⊤, (ii) aligns it to the corresponding prefix of the global reference basis via orthogonal Procrustes, (iii) reconstructs a basis-consistent local update, (iv) aggregates with weights proportional to p_i · c_i (data-size weight times consistency score), (v) refactorizes with an exact SVD of the small coefficient matrix C_g^t ∈ ℝ^{m×r_max} rather than the full m×n update, and (vi) applies gauge fixing before rank-prefix dispatch.

3. Theory tying consistency to error. Theorem 3.3 gives min_T ||Q_i^tT − Q_ref,i^{t−1}||_F^2 = 2r_i(1−c_i^t) and bounds the reconstruction residual by s||B̄_i^t||_2√(2r_i(1−c_i^t)). Theorem 3.5 shows consistency-aware weighting tightens an aggregate reconstruction bound versus pure data-size weighting, and Theorem 3.8 places the residual into a standard smooth non-convex convergence bound.

4. Empirical validation. Experiments against six heterogeneous-rank federated LoRA baselines on GLUE, Dolly-15K/Alpaca/WizardLM, and Natural Instructions v2, plus per-rank analysis, rank-heterogeneity regimes, basis diagnostics, sensitivity analysis, system-efficiency measurements, and ablations.

Main Findings

• Basis inconsistency is real and worsens with rank heterogeneity. On LLaMA-2-7B with N=8 clients on Dolly-15K using FlexLoRA as backbone, at round t=10 the consistency score distribution shifts downward as heterogeneity grows; even in the homogeneous case the median sits well below 1, and under high heterogeneity a substantial mass of c_i^t values falls below 0.5.

• Consistency is only weakly related to data size. Under high heterogeneity, c_i^t versus the data-size weight p_i shows |ρ| < 0.2, and the clients with the largest p_i are not necessarily those with the largest c_i^t. Pure data-size weighting can therefore amplify poorly aligned clients.

• Best overall accuracy across task families. FedAbA reports 90.3 MRPC F1, 63.6 CoLA MCC, 94.2 QNLI accuracy, 83.5 RTE accuracy, 37.1 (Wizard) / 47.8 (Dolly-15K) / 37.8 (Alpaca) MMLU, 4.32 MT-Bench on Alpaca, and 58.4 ROUGE-L / 66.1 ROUGE-1 on Natural Instructions v2. The next-best baseline, Fed-PLoRA, reports 89.5, 63.1, 94.0, 81.8, 36.6, 47.2, 37.3, 4.23, 58.1, and 65.8 respectively.

• Largest gains on the lowest ranks. Per-rank evaluation on Dolly-15K gives FedAbA 31.2 at r=4, 32.0 at r=8, 32.5 at r=16, 33.0 at r=32, 33.4 at r=64, overall 32.4; on Natural Instructions v2 it gives 57.4, 58.6, 59.0, 60.6, overall 58.4. The paper attributes this to low-rank clients receiving only short global prefixes and being most exposed to cross-round basis drift.

• Margins widen with heterogeneity. FedAbA's NLU average is 80.97 (homogeneous [8]×10), 81.72 (mild {4,8,16}), and 82.90 (high {4,8,16,32,64}); NLG MMLU average is 39.30 / 39.83 / 40.90; NI ROUGE average is 61.60 / 61.90 / 62.25. The best competitor under high heterogeneity is ECLoRA at 82.20 (NLU) and Fed-PLoRA at 40.37 (NLG) and 61.95 (NI).

• Better basis diagnostics. Under high rank heterogeneity, FedAbA yields higher alignment scores and a lower sign/order flip rate across consecutive global bases than the compared baselines.

• Favorable system cost. FedAbA communicates 520 MB per round for N=10, has effective rank 64, takes 2.0 s aggregation time, and uses 190 MB server memory. For comparison, FlexLoRA takes 130 s and 2.15 GB, ECLoRA 4.0 s and 4.3 GB, FLoRA 2.86 GB communication with effective rank 248, and LoRA-A²/Fed-PLoRA 931 MB communication.

• Every component contributes. Ablation on NLU average / NLG MMLU / MT-Bench / NI ROUGE-L / ROUGE-1: removing basis consistency gives 82.02 / 39.95 / 4.15 / 57.8 / 65.6; rotation-only 82.07 / 40.00 / 4.16 / 57.9 / 65.6; removing Procrustes 82.32 / 40.20 / 4.20 / 58.0 / 65.8; removing c_i^t weighting 82.42 / 40.40 / 4.25 / 58.1 / 65.9; removing data size 82.62 / 40.55 / 4.27 / 58.2 / 66.0; removing gauge fixing 82.52 / 40.50 / 4.26 / 58.2 / 65.9; full FedAbA 82.90 / 40.90 / 4.32 / 58.4 / 66.1.

• Stable across repeated runs. Three independent matched runs comparing FedAbA and Fed-PLoRA on six selected metrics report a higher mean for FedAbA on all six (Appendix F.1); on the 20% of Natural Instructions clients with the lowest average consistency, FedAbA also achieves higher average ROUGE-L and ROUGE-1 (Appendix F.2).

• The method operates within a fixed span. Remark 2.1 shows span(Q_g^t) = span(Q_g^{t−1}), so the SVD refresh plus gauge fixing never moves the global basis out of the warm-start span; warm start uses 5 rounds of standard federated LoRA aggregation plus a rank-r_max truncated SVD.

Methodology in Plain English

The server keeps a "reference direction set" — the transpose of the previous round's global right factor A_g^{t−1}, which is orthonormal by construction. Each client sends its two LoRA matrices. Instead of averaging the composed matrices directly, the server:

1. Separates direction from magnitude. A thin QR of the client's right factor splits it into an orthonormal local basis Q_i^t (directions) and a small coefficient matrix C_i^t (magnitudes), so that ΔW_i^t = sB̄_i^t(Q_i^t)^⊤.

2. Aligns the basis. It solves an orthogonal Procrustes problem to find the rotation T_i^t that best maps the local basis onto the client's reference prefix. The same computation yields the consistency score c_i^t, the average cosine of the principal angles between the two bases.

3. Rebuilds the update on the shared basis. Crucially, it does not just rotate the factors — that would change nothing after composition. It substitutes the reference prefix for the local basis, producing ΔW̃_i^t = sB̄_i^tT_i^t(Q_ref,i^{t−1})^⊤, keeping all r_i coefficient columns. The leftover mismatch is a residual bounded by s||B̄_i^t||_2√(2r_i(1−c_i^t)).

4. Aggregates with a sanity check on alignment. Weights are α_i^t = n_i c_i^t / Σ_j n_j c_j^t, so a data-rich but badly aligned client no longer dominates. If all active clients have c_i^t = 0, it falls back to plain data-size weights.

5. Refactorizes cheaply and tidily. Because the aggregate equals sC_g^t(Q_g^{t−1})^⊤ with C_g^t ∈ ℝ^{m×r_max}, the server SVDs only that small matrix instead of a full m×n one. A gauge-fixing step then matches near-degenerate direction blocks back to the previous basis by maximum absolute cosine similarity, choosing signs to keep matched cosines nonnegative — a relabeling that leaves the composed update unchanged but keeps the prefixes sent to clients stable.

Client-side training and communication are unchanged; all extra work is on the server. Per-client cost is O(nr_i²) for the QR, O(nr_i² + r_i³) for Procrustes, O(mr_i²) for coefficient transforms, O(|S^t| m r_max) for aggregation, and O(m r_max² + r_max³) for the final SVD; communication stays O(r_i(m+n)) per client.

Why This Matters

Research impact. The paper reframes heterogeneous-rank federated LoRA from a bookkeeping problem (pad, stack, or truncate ranks) to a geometric one: which basis are we averaging in? The distinction between rotating factors and reconstructing updates is a clean, provable insight, and the consistency score gives practitioners a cheap diagnostic for prefix stability that does not require access to raw data.

Real-world applications (plausible deployments, not claims evaluated in the paper):

• Cross-device personalization of on-device language assistants, where phones, tablets, and laptops can only afford ranks from 4 to 64.
• Hospital or clinic consortia fine-tuning a shared clinical language model without pooling patient records.
• Enterprise multi-tenant adaptation of a foundation model, where business units have different GPU budgets and cannot share proprietary text.
• Edge fleets in manufacturing or retail running local instruction-tuned models over sensor and transaction logs.

Industry relevance. The method requires no client-side changes, keeps per-round communication at the same 520 MB level as HetLoRA and FlexLoRA, and avoids FlexLoRA's 130 s aggregation time and 2.15 GB server memory, or ECLoRA's 4.3 GB. That makes it a drop-in server upgrade for existing federated LoRA deployments, and the exact coefficient-space refactorization avoids materializing full weight-sized matrices on the server.

Future Directions

• Escape the fixed warm-start span. Remark 2.1 establishes that the global basis never leaves span(Q_g^{t−1}), so FedAbA cannot discover genuinely new directions after warm start. A method that periodically expands the span while preserving prefix stability is an open problem.
• Adaptive rank assignment. Ranks are fixed per client in the experiments ([4,4,8,8,16,16,32,32,64,64] for NLU and NLG-style; balanced over 4, 8, 16, 32 for NI). Whether the consistency score could drive rank allocation across rounds is unexplored.
• Relaxing the idealized analysis. Theorem 3.8 rests on Assumption 3.7 (unbiased stochastic gradients, bounded variance) and Theorem 3.5 on Assumption 3.4 (a uniform coefficient-radius bound κ_t). Behavior under partial participation beyond the reported setting, client drift, and asynchronous updates is not analyzed.
• Combining with compression, privacy, and security. The paper does not study quantization or sparsification of the coefficient matrices, nor secure aggregation or differential privacy, all of which are natural companions for real federated LLM training.

Target Audience

Researchers and engineers working on federated learning, parameter-efficient fine-tuning, or distributed LLM adaptation — particularly those deploying LoRA across devices or organizations with unequal resources. It suits readers comfortable with matrix factorization and SVD-based analysis; practitioners mainly interested in deployment can read the method and experiment sections, while the theory in Section 3 and the appendices is aimed at those who want the formal justification for basis-consistent aggregation and consistency-aware weighting.

Authors’ abstract

Federated fine-tuning of large language models (LLMs) under heterogeneous client capacities is increasingly important, and low-rank adaptation (LoRA) with heterogeneous ranks makes it practical. Existing server-side aggregators compose local matrix updates $ΔW_i$ from client-specific LoRA factors $(A_i,B_i)$ before aggregating them. Although these matrix updates lie in the same weight space $\mathbb{R}^{m\times n}$ and can be averaged directly, the LoRA factors that generate them may expressed client-specific local basis. In the typical compose--aggregate--factorize--dispatch loop, this basis inconsistency can make the rank-$r_i$ reference prefix dispatched to lower-rank clients unstable after server-side refactorization. In this work, we propose \textbf{FedAbA}, a basis-consistent federated LoRA aggregation framework for heterogeneous ranks. FedAbA first extracts each client's local basis, aligns it to the corresponding prefix of the global reference basis, and constructs a basis-consistent reconstruction of each local matrix update before aggregation. It then aggregates these reconstructions with weights that combine client data size and consistency score, followed by exact coefficient-space SVD refactorization and gauge fixing before rank-prefix dispatch. Our theoretical analysis clarifies the role of basis consistency and supports the proposed consistency-aware aggregation mechanism. Extensive experiments show that FedAbA outperforms representative federated LoRA baselines, with the largest gains under the strongest rank heterogeneity.