PaperlayerView exact revision on arXiv

Source-linked AI summary

Selective Aggregation for Low-Rank Adaptation in Federated Learning

Pengxin Guo, Shuang Zeng, Yanran Wang, Huijie Fan, Feifei Wang, Liangqiong Qu

arXiv:2410.01463v3cs.LG

TL;DR

Federated fine-tuning of pretrained language models faces data-sharing restrictions and non-IID client data, while directly aggregating LoRA matrices can yield a different update from the ideal one. The paper analyzes the distinct roles of A and B, proposes sharing only A in FedSA-LoRA, and reports that its methods outperform other approaches across natural language understanding and generation tasks.

  • Problem

    Federated learning supports collaborative training without external sharing of personal data, but directly aggregating LoRA’s A and B matrices can produce an update different from the ideal update.

    Source: §1 · p. 2

  • Method

    The paper analyzes the learned A and B matrices and proposes FedSA-LoRA, which trains both but shares only A for server aggregation, extending the approach to rsLoRA and VeRA.

    Source: §1 · p. 2

  • Results

    The authors report that FedSA-LoRA, FedSA-rsLoRA, and FedSA-VeRA outperform other methods in extensive experiments on natural language understanding and generation tasks.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The authors present selective sharing of A, while keeping B local, as a general paradigm for combining LoRA variants with federated learning.

    Source: §1 · p. 3

    The convergence analysis assumes that all client loss functions are L-smooth.

    Source: §4.2 · p. 6

Abstract

from arXiv · show

We investigate LoRA in federated learning through the lens of the asymmetry analysis of the learned $A$ and $B$ matrices. In doing so, we uncover that $A$ matrices are responsible for learning general knowledge, while $B$ matrices focus on capturing client-specific knowledge. Based on this finding, we introduce Federated Share-A Low-Rank Adaptation (FedSA-LoRA), which employs two low-rank trainable matrices $A$ and $B$ to model the weight update, but only $A$ matrices are shared with the server for aggregation. Moreover, we delve into the relationship between the learned $A$ and $B$ matrices in other LoRA variants, such as rsLoRA and VeRA, revealing a consistent pattern. Consequently, we extend our FedSA-LoRA method to these LoRA variants, resulting in FedSA-rsLoRA and FedSA-VeRA. In this way, we establish a general paradigm for integrating LoRA with FL, offering guidance for future work on subsequent LoRA variants combined with FL. Extensive experimental results on natural language understanding and generation tasks demonstrate the effectiveness of the proposed method. Our code is available at https://github.com/Pengxin-Guo/FedSA-LoRA.

1 INTRODUCTION

The introduction frames federated fine-tuning as a response to restricted data sharing and motivates moving beyond uniform client updates. It identifies distinct roles for LoRA’s A and B matrices and introduces FedSA-LoRA, which shares only A for aggregation.

1 INTRODUCTION

  • Directly aggregating both LoRA matrices can produce a different update from the ideal update, while fixing A can impair LoRA’s learning ability.The introduction also motivates personalized federated learning because uniform updates are not optimal, especially under non-IID data.Source: §1 · p. 2
  • A matrices learn general knowledge, whereas B matrices capture client-specific knowledge; the learned A matrices are more similar across clients, while B similarity decreases with greater data heterogeneity.The authors use this distinction to motivate selective aggregation.Source: §1 · p. 2
  • FedSA-LoRA trains both low-rank matrices but shares only A with the server for aggregation, keeping B local.The method permits different client model updates rather than requiring a uniform update.Source: §1 · pp. 2–3
  • The authors extend the approach to rsLoRA and VeRA as FedSA-rsLoRA and FedSA-VeRA, reporting a similar A/B pattern in those variants.They present these extensions as a general paradigm for combining LoRA variants with federated learning.Source: §1 · p. 3
  • The authors report that FedSA-LoRA, FedSA-rsLoRA, and FedSA-VeRA outperform other methods in extensive experiments.The introduction describes results across natural language understanding and generation tasks.Source: §1 · p. 3

2 RELATED WORK

The related work covers federated fine-tuning under non-IID data and resource constraints, then reviews LoRA-based approaches to federated learning. It distinguishes FedSA-LoRA from FFA-LoRA by keeping both matrices trainable while sharing only A.

2.1 FEDERATED LEARNING

  • Federated fine-tuning faces non-IID performance discrepancies and substantial communication, computation, and memory demands; parameter-efficient techniques, including LoRA, are explored for resource-limited settings.LoRA modules can be merged with frozen weights at deployment without inference latency.Source: §2.1 · pp. 3–4

2.2 LORA IN FEDERATED LEARNING

  • Federated LoRA methods include FedLoRA, DP-LoRA, FedDPA, and FDLoRA, with some using personalized and global LoRA modules to represent distinct knowledge.Other work handles heterogeneous ranks through zero-padding, replication, weighted aggregation, or stacking.Source: §2.2 · p. 4
  • FFA-LoRA freezes A and aggregates trainable B, whereas FedSA-LoRA keeps both matrices trainable and shares only A for aggregation.The related-work passage notes that fixing matrices can impair LoRA’s learning ability and yield suboptimal performance.Source: §2.2 · p. 4

3 MOTIVATING EXAMPLE

The motivating example uses a low-rank update to analyze the roles of LoRA’s A and B matrices. Its analysis and experiments associate A with general knowledge and B with client-specific knowledge, including across other LoRA variants.

3 MOTIVATING EXAMPLE

  • LoRA represents a weight update as BA while keeping pretrained weights W0 frozen; A and B are trainable low-rank matrices, with B initialized to zero and A randomly initialized.The fine-tuned weight is W0 + BA.Source: §3 · p. 4
  • The lemma-based analysis says the optimal A is independent of the input data distribution, while the optimal B depends on it, supporting distinct general-knowledge and client-specific roles.The example considers fine-tuning either factor while fixing the other.Source: §3 · p. 5
  • In RoBERTa-large on RTE with three clients, A matrices were more similar across clients than B matrices, and B similarity decreased as data heterogeneity increased.The experiments used IID data and Dirichlet non-IID partitions with α = 1 and α = 0.5; the authors also report that A was updated from initialization.Source: §3 · p. 5
  • The A/B similarity pattern also appeared in rsLoRA and VeRA; the authors present this extension as a general phenomenon for LoRA with federated learning.They note that prior asymmetry analyses did not provide this generalization across the variants and non-IID settings considered here.Source: §3 · pp. 5–6

4 OUR METHOD

FedSA-LoRA aggregates only A while clients retain B locally, separating general from client-specific knowledge. The paper extends this design to rsLoRA and VeRA and provides a convergence analysis.

4.1 FEDERATED SHARE-A LOW RANK ADAPTATION

  • FedSA-LoRA trains A and B to represent weight updates, shares A for server aggregation, and keeps client-specific B matrices local.The local model is W0 + BA; the server averages and broadcasts A between training rounds.Source: §4.1 · p. 6
  • FedSA-rsLoRA retains the FedSA-LoRA design with a different scaling factor, while FedSA-VeRA aggregates Ad and trains Bb locally.In VeRA, A and B are fixed shared matrices adapted by trainable scaling vectors Ad and Bb.Source: §4.1 · p. 6

4.2 CONVERGENCE ANALYSIS

  • The convergence assumptions include L-smooth client losses, bounded expected squared stochastic-gradient norms, and nonzero eigenvalues for A and B.The paper identifies the smoothness and bounded-second-moment assumptions as common in FL work.Source: §4.2 · pp. 6–7
  • The analysis gives a convergence rate toward a stationary solution comparable to traditional FedAvg under smooth, non-convex conditions.The theorem specifies assumptions and a learning rate; E denotes local training iterations between communication rounds.Source: §4.2 · p. 7

5 EXPERIMENTS

FedSA variants consistently outperform other methods across the reported natural language understanding tasks, while additional analyses examine efficiency, data heterogeneity, client counts, LoRA rank, and generation tasks.

5.1 NATURAL LANGUAGE UNDERSTANDING

  • FedSA-LoRA, FedSA-rsLoRA, and FedSA-VeRA consistently outperform other methods across all GLUE tasks.Table 1 reports accuracy across three runs with mean and standard deviation.Source: §5.1 · p. 8, Table 1 · p. 8
  • FedSA-LoRA reaches the RTE target in 91 rounds versus 229 for FFA-LoRA, despite slightly higher per-round computation cost (22s versus 20s).The authors report that FedSA-LoRA requires the smallest communication cost to reach target performance.Source: §5.1 · p. 8, Table 2 · p. 9

5.2.1 EFFECT OF DATA HETEROGENEITY

  • FedSA-LoRA's accuracy gains over LoRA increase with heterogeneity, reaching 1.84% on QNLI and 1.4% on MNLI-m under severe non-IID conditions.The authors link the pattern to decreasing similarity between clients' B matrices as heterogeneity increases.Source: §5.2.1 · p. 9, Table 3 · p. 9

5.2.2 EFFECT OF NUMBER OF CLIENTS

  • FedSA-LoRA outperforms other methods with both three clients and larger client counts from 10 to 100.The 100-client setting uses client sampling at rate 0.3.Source: §5.2.2 · p. 9, Table 4 · p. 10

5.2.3 EFFECT OF LORA RANK

  • FedSA-LoRA outperforms other methods across LoRA ranks 2, 4, 8, and 16 on QNLI, SST2, and MNLI-m.The authors describe its performance as adaptable and robust across these rank settings.Source: §5.2.3 · p. 10, Table 5 · p. 10

5.3 NATURAL LANGUAGE GENERATION

  • FedSA-LoRA scores 46.63 on GSM8K versus 46.23 for LoRA and 46.32 for FFA-LoRA, and 59.66 on CodeSearchNet versus 58.34 and 58.57.The GSM8K evaluation uses three IID clients; CodeSearchNet uses a default non-IID partition.Source: §5.3 · p. 10

6 CONCLUSION

The conclusion assigns A matrices the role of learning general knowledge and B matrices the role of capturing client-specific knowledge. It presents FedSA-LoRA and extensions to rsLoRA and VeRA as a general paradigm for combining LoRA with federated learning.

  • A matrices learn general knowledge, while B matrices capture client-specific knowledge when LoRA is combined with federated learning.Source: §6 · p. 10
  • FedSA-LoRA trains both low-rank matrices but shares only A matrices with the server for aggregation.Source: §6 · p. 10
  • The method extends to rsLoRA and VeRA as FedSA-rsLoRA and FedSA-VeRA, based on a consistent A/B relationship across variants.Source: §6 · p. 10
  • The paper presents this approach as a general paradigm for integrating LoRA with federated learning and guiding future LoRA variants.Source: §6 · p. 10

A.1 PROOF OF LEMMA 1

The proof examines the loss and gradients under two cases: updating B with A fixed, and updating A with B fixed. In each case, it sets the corresponding gradient to zero to characterize the optimal matrix.

  • For the B-update case, the proof fixes A = Q and considers the loss in Eq. (3).Source: §A.1 · p. 16
  • It derives the gradient with respect to B and sets it to zero to obtain the condition for optimal B∗.Source: §A.1 · p. 16
  • For the A-update case, the proof fixes B = U, considers the loss, and derives the gradient with respect to A.Source: §A.1 · p. 16
  • It sets the A-gradient to zero to obtain the condition for optimal A∗.Source: §A.1 · p. 16

A.2 PROOF OF THEOREM 1

The proof describes a client’s one-step update, including the communication step, and introduces notation for the sub-step parameters and stochastic gradient. It then applies smoothness and expectation arguments under stated assumptions to derive bounds across steps.

  • The proof frames the proposed method as a one-step update for client i when the next step is a global synchronization step.Source: §A.2 · p. 17
  • It introduces notation for parameters at intermediate sub-steps and for the parameter obtained after communication, and defines the client’s learning rate and stochastic gradient.Source: §A.2 · p. 17
  • The stochastic gradient uses data uniformly chosen from client i’s local dataset at step t.Source: §A.2 · p. 17
  • The proof applies smoothness and total-expectation arguments and invokes positive constants cA and cB in its assumptions.Source: §A.2 · pp. 17–18
  • It combines intermediate inequalities, repeatedly applies Eq. (22), and divides the accumulated bound by T.Source: §A.2 · pp. 19–20

A.3 FURTHER EXPLANATIONS ABOUT AGGREGATION ERRORS.

The section contrasts the ideal average of client update products with the update obtainable when LoRA’s trainable factors are averaged separately. The resulting product differs from the ideal update and introduces aggregation errors; averaging the ideal weight update also cannot be followed by continued federated training because it cannot be decomposed back into A and B.

  • The ideal server update averages each client’s product B_iA_i across the m clients.Source: §A.3 · p. 21
  • Because LoRA trains A_i and B_i rather than W_i, practical aggregation separately averages the factors and then combines them.Source: §A.3 · p. 21
  • The resulting factor averages include cross-client pairings and differ from the ideal average of client products, introducing aggregation errors.Source: §A.3 · p. 21
  • Averaging the ideal weight updates cannot support further federated training because the averaged W cannot be decomposed back into A and B.Source: §A.3 · p. 21

A.4 HYPERPARAMETERS

The appendix identifies Table 6 as the GLUE learning-rate table for LoRA-based methods and describes the optimizer and learning-rate strategy used for VeRA-based methods.

A.4 HYPERPARAMETERS

  • VeRA-based methods used AdamW with separate learning rates for the classification head and adapted layers after tested SGD rates performed significantly worse than LoRA-based methods.The best of VeRA, FFA-VeRA, and FedSA-VeRA reached 53.73% on MNLI-m, compared with 90.20% for LoRA-based methods.Source: §A.4 · p. 21
  • Table 6 lists learning rates used for LoRA-based methods on the GLUE benchmark.Source: Table 6 · p. 21

A.5 EXAMPLES OF GENERATED ANSWER FOR GSM8K DATASETS

The appendix reports a GSM8K example in which FedSA-LoRA answers accurately where LoRA and FFA-LoRA make reasoning errors, then compares learned matrices and tests A-matrix compression.

A.5 EXAMPLES OF GENERATED ANSWER FOR GSM8K DATASETS

  • FedSA-LoRA answers the GSM8K example accurately, while LoRA and FFA-LoRA make reasoning errors.Table 9 reports GSM8K performance and example answers generated by each method.Source: §A.5 · p. 21, Table 9 · p. 22

A.6.1 LEARNED AND INITIALIZED LORA MATRICES COMPARISON

A.6.2 LEARNED RSLORA MATRICES COMPARISON

A.6.3 LEARNED VERA MATRICES COMPARISON

  • VeRA fixes shared A and B matrices and adapts trainable scaling vectors; Ad is more similar across clients than Bb, whose similarity falls as heterogeneity increases.Source: §A.6.3 · p. 24, Figure 6 · p. 24

A.7 FURTHER RESULTS ON COMPRESSING A MATRICES

  • FetchSGD compresses around 50% of A matrices, with performance comparable to original FedSA-LoRA while reducing communication overhead.Table 10 reports time and space costs on RTE and QNLI; its caption marks compressed methods with ‡.Source: §A.7 · p. 25, Table 10 · p. 25

Related papers

Loading 2410.01463v3…