Source-linked AI summary

Improving LoRA in Privacy-preserving Federated Learning

Youbang Sun, Zitao Li, Yaliang Li, Bolin Ding

arXiv:2403.12313v1cs.LGcs.CRcs.DC

TL;DR

The paper addresses instability in LoRA-based privacy-preserving federated learning caused by heterogeneous local optimization, DP-noise amplification, and scaling-factor sensitivity. It proposes FFA-LoRA, which freezes the randomly initialized non-zero matrix and trains only the zero-initialized matrix. Across the reported experiments, FFA-LoRA consistently outperforms LoRA with lower computational and communication costs.

  • Problem

    LoRA’s joint local optimization and separate server aggregation create discordance in privacy-preserving FL, alongside DP-noise amplification and sensitivity to α.

  • Method

    FFA-LoRA freezes the non-zero initialized low-rank matrix and updates and aggregates only the zero-initialized matrix.

  • Results

    FFA-LoRA consistently outperforms LoRA across the paper’s federated experiments and behaves better in both i.i.d. and non-i.i.d. settings.

  • Takeaways & Limitations

    FFA-LoRA provides more consistent performance with improved computational efficiency while preserving LoRA’s benefits in privacy-preserving FL.

  • Takeaways & Limitations

    The paper identifies the initialization of fixed A as an open performance question and discusses it in an appendix ablation.

Abstract

from arXiv · show

Low-rank adaptation (LoRA) is one of the most popular task-specific parameter-efficient fine-tuning (PEFT) methods on pre-trained language models for its good performance and computational efficiency. LoRA injects a product of two trainable rank decomposition matrices over the top of each frozen pre-trained model module. However, when applied in the setting of privacy-preserving federated learning (FL), LoRA may become unstable due to the following facts: 1) the effects of data heterogeneity and multi-step local updates are non-negligible, 2) additive noise enforced on updating gradients to guarantee differential privacy (DP) can be amplified and 3) the final performance is susceptible to hyper-parameters. A key factor leading to these phenomena is the discordance between jointly optimizing the two low-rank matrices by local clients and separately aggregating them by the central server. Thus, this paper proposes an efficient and effective version of LoRA, Federated Freeze A LoRA (FFA-LoRA), to alleviate these challenges and further halve the communication cost of federated fine-tuning LLMs. The core idea of FFA-LoRA is to fix the randomly initialized non-zero matrices and only fine-tune the zero-initialized matrices. Compared to LoRA, FFA-LoRA is motivated by practical and theoretical benefits in privacy-preserved FL. Our experiments demonstrate that FFA-LoRA provides more consistent performance with better computational efficiency over vanilla LoRA in various FL tasks.

1 INTRODUCTION

The paper studies LoRA for privacy-preserving federated fine-tuning, where data heterogeneity, private optimization, and hyper-parameter sensitivity create challenges. It proposes FFA-LoRA and reports consistent improvements over LoRA.

  • LoRA is attractive for federated fine-tuning because it reduces communication and local computation for parties with insufficient or private task-specific data.
  • Directly combining LoRA with FedAvg and differential privacy can be suboptimal because noisy gradients and non-i.i.d. client data interfere with optimization.
  • The paper identifies three discordances: separate aggregation of jointly optimized matrices, DP-noise amplification, and strong dependence on LoRA’s scaling factor α.
  • FFA-LoRA freezes the non-zero initialized low-rank matrices and updates only the zero-initialized matrices, halving trainable parameters and associated FL communication and computation.
  • Experiments are designed to evaluate FFA-LoRA across tasks, hyper-parameters, and privacy protection levels, with the reported outcome that it consistently outperforms LoRA.

2 BACKGROUND AND RELATED WORKS

The background motivates parameter-efficient federated fine-tuning for large, distributed, and private datasets. It introduces PEFT, federated learning, differential privacy, and DP-SGD as the relevant framework and mechanism.

  • PEFT adds a small set of trainable parameters while keeping most pre-trained parameters frozen, encoding the task-specific update in a lower-dimensional representation.
  • Federated learning lets multiple parties collaboratively fine-tune LLMs when each party has inadequate domain-specific data and cannot directly share its dataset.
  • PEFT methods in FL reduce communication and local training burdens, but data heterogeneity remains a challenge for PEFT algorithms.
  • Differential privacy limits how much any single record affects output distributions, with smaller (ε, δ) values indicating stronger protection.
  • DP-SGD clips per-sample gradients and adds Gaussian noise before using the privatized gradient to update the model.
  • The paper adopts client-to-server DP, so shared local model updates satisfy DP and the final model inherits global privacy guarantees through composition and post-processing.

3 LORA IN PRIVACY-PRESERVING FL

LoRA’s two-factor parameterization becomes discordant with privacy-preserving FL because clients jointly optimize its factors while servers aggregate them separately. This creates aggregation error, amplified DP noise, and α sensitivity.

  • LoRA updates frozen weights through two low-rank matrices A and B, reducing trainable parameters when their rank r is much smaller than the layer dimensions.
  • Vanilla LoRA jointly optimizes A and B locally but averages them separately globally, so the product of averaged factors need not equal the averaged client updates.
  • This mismatch becomes more significant with more local update steps and more different client datasets, linking LoRA to client drift in heterogeneous FL.
  • DP-SGD noise is cascaded through LoRA’s matrix multiplication, introducing the non-Gaussian product term ξBξA and additional convergence difficulty.
  • Approximately 3 times more noise is produced by LoRA than full-model fine-tuning for 1000 rounds, 10 local steps, SST-2, batch size 200, ε = 6, and δ = 1e−5.
  • Increasing α can improve performance on complex tasks but also increases instability, noise impact, and the costly hyper-parameter search burden.

4 A SIMPLE RECEIPT: FFA-LORA

FFA-LoRA fixes the randomly initialized matrix A and trains only the zero-initialized matrix B. The resulting linearized update is more compatible with FL aggregation and DP, while reducing efficiency and α-tuning costs.

  • FFA-LoRA uses W = W0 + BA0, keeping the pre-trained weight W0 and random Gaussian matrix A0 frozen while training B from zero initialization.
  • Because every client shares the same fixed A0, FFA-LoRA’s updates are compatible with FedAvg and avoid LoRA’s low-rank aggregation error term.
  • FFA-LoRA avoids DP-noise amplification because noise enters only linearly through the trained B matrix, without LoRA’s product-noise term.
  • FFA-LoRA does not rely on α and is equivalent to LoRA with α = ∞ under the stated trajectory connection.
  • Fixing A0 effectively halves trainable parameters, reducing computation and communication while allowing α to be fixed and narrowing hyper-parameter search.
  • The paper concludes that FFA-LoRA preserves LoRA’s benefits while avoiding its identified shortcomings in privacy-preserving FL.

5 EXPERIMENTS

Experiments compare FFA-LoRA and LoRA across language understanding, language generation, privacy, heterogeneity, rank, and scaling settings. FFA-LoRA generally performs more consistently and efficiently, with a reported GSM-8K advantage over LoRA.

  • Experimental setup: Experiments evaluate FFA-LoRA and LoRA on RoBERTa language-understanding tasks and LLaMA language-generation tasks.The language-understanding tasks are MNLI, SST2, QNLI, and QQP; generation uses GSM-8K.
  • Experimental setup: The experiments use a 3-client cross-silo federated setting with deliberately heterogeneous client data and 1000 communication rounds.The setup keeps batch size B = 200 and fixes the comparison conditions across clients.
  • Overall comparison: DP significantly degrades performance for both methods, yet FFA-LoRA performs better with and without privacy, especially on strongly heterogeneous MNLI.The authors interpret the largest MNLI gap as evidence that FFA-LoRA is more suitable under strong heterogeneity.
  • Data heterogeneity: FFA-LoRA performs better than LoRA in both i.i.d. and non-i.i.d. settings, while performance is similar in the privacy-free setting.The comparison is reported for four tasks under homogeneous and heterogeneous data distributions.
  • Rank and privacy budgets: Across rank selections, FFA-LoRA performs better on most tasks regardless of trainable parameter number.The study evaluates ranks r ∈ {2, 4, 8, 16} and notes that increasing rank does not necessarily increase gradient information.
  • Rank and privacy budgets: Under stronger privacy requirements, LoRA becomes more sensitive to rank, whereas FFA-LoRA maintains relatively stable performance across rank selections.This comparison uses privacy budgets ϵ ∈ {6, 3, 1} and ranks r ∈ {2, 4, 8, 16}.
  • Hyper-parameter studies: The experiments include scaling-factor studies and discuss how initialization of the fixed A matrix affects FFA-LoRA performance.The paper states that scaling does not affect FFA-LoRA’s overall performance and refers initialization analysis to Appendix A.8.
  • Language generation: 17.12% accuracy on GSM-8K exceeds LoRA’s 15.68% best performance.The passage also reports 15.31% from prior work as a separate LoRA result.

6 CONCLUSION

The paper analyzes LoRA in privacy-preserving federated learning and proposes FFA-LoRA as a theoretically motivated, empirically verified, and more computationally efficient modification.

  • Conclusion: FFA-LoRA improves LoRA for privacy-preserving federated learning while reducing the computational burden of federated LLM fine-tuning.The paper also identifies future directions involving more parameter-efficient algorithms and alternative matrix initialization methods.

A.1 SMOOTHNESS ANALYSIS

With A fixed and bounded, the composed loss remains Lipschitz smooth in the trainable matrix B; when both low-rank matrices are trainable, smoothness is not guaranteed.

  • Fixing A with ∥A∥≤C preserves Lipschitz smoothness in B with factor LC2.
  • When both A and B are trainable, the composed loss has no Lipschitz smoothness guarantee.
  • The analysis assumes smoothness is measured using the matrix Frobenius norm.
  • For W(A, B) = W0 + BA, the gradient with respect to B is obtained from the W-gradient as ∇BF = ∇W FAT.
  • The corresponding gradient with respect to A is BT∇WF, allowing the smoothness properties to be analyzed for both matrices.

1. For property 1, we know that for any given B1, B2,

A counterexample shows that a loss can be Lipschitz smooth in the original model parameter W while failing to be smooth in the jointly parameterized low-rank variables.

  • A sequence with Ak = Bk = kId demonstrates unbounded gradient variation in the stacked variable x = [A, B].
  • Although F(W) is 1-Lipschitz smooth, the composed function is not smooth with respect to x.

A.2 PROOF FOR THEOREM 1

Theorem 1 compares LoRA and FFA-LoRA updates under the same initial model condition, showing that FFA-LoRA updates only B and avoids dependence on its scaling factor.

  • The theorem uses W = W0 + αBαAα and LoRA’s non-zero A initialization to identify Aα with the fixed matrix A.
  • For FFA-LoRA, different αFFA values do not affect the update dynamics, so the proof may consider αFFA = 1.
  • FFA-LoRA updates and aggregates only B while keeping A fixed.
  • The update comparison is extended across local iterations by induction.
  • Changing αLoRA changes LoRA’s updates because LoRA does not satisfy the theorem’s conditions.

A.3 LLAMA EXPERIMENTS AND DETAILS

On LLaMA with GSM-8K, FFA-LoRA achieves higher accuracy than LoRA under the same hyper-parameters and performs better on longer, more complex questions.

  • 17.12% accuracy on GSM-8K exceeds LoRA’s best 15.68% performance, or 15.31% previously reported.
  • FFA-LoRA’s generated answers demonstrate better performance on longer questions with more complex contexts.

A.4 A MOTIVATION FOR FURTHER REDUCING TRAINABLE PARAMETERS

The paper examines whether adapter designs can reduce trainable parameters further while retaining meaningful learning and privacy guarantees. It presents FFA-LoRA and contrasts it with more aggressive parameter reduction through QVP adapters and experiments across settings.

  • FFA-LoRA: FFA-LoRA freezes randomly initialized parameters and trains only zero-initialized parameters, reducing trainable parameters relative to LoRA.The paper motivates this design as theoretically beneficial and reports lower communication and computational cost in federated learning.
  • QVP adapters: For QVP adapters, randomly initialized Q0 and P0 remain frozen while the zero-initialized V matrix is trainable.The construction uses W0 as the pre-trained weight and compares QVP with LoRA and FFA-LoRA.
  • QVP adapters: QVP adapters underperform LoRA and FFA-LoRA at matched parameter budgets but can reduce trainable parameters further when the budget is extremely constrained.The paper identifies local private training on mobile devices as a potential use case.
  • Differential privacy: FFA-LoRA’s privacy guarantee follows from locally applied DP-SGD, parallel composition across clients, server-side averaging as post-processing, and composition across rounds.The resulting mechanism satisfies (ϵ, δ)-DP under the stated accountant and update conditions.
  • Scaling-factor experiments: Across scaling factors α, FFA-LoRA remains consistent when the learning rate scales with α, whereas LoRA can degrade drastically and requires arduous tuning.The paper reports that LoRA can still converge with an adequate learning rate, but finding it for a given α is generally difficult.
  • Experimental comparisons: In centralized experiments, FFA-LoRA and vanilla LoRA show no significant performance discrepancy.This result is reported without considering DP and FL.
  • Experimental comparisons: In iid federated experiments, FFA-LoRA shows better convergence and fewer training fluctuations than LoRA.The paper states that this pattern also appears in language-related tasks and is not limited to language tasks.
  • Initialization: Orthogonal initialization of matrix A performs slightly better than the existing initialization in initial results, but the performance gap is not definitive.The paper also considers Kaiming initialization and the top r singular vectors of W0.
Loading 2403.12313v1…