Source-linked AI summary

SLoRA: Federated Parameter Efficient Fine-Tuning of Language Models

Sara Babakniya, Ahmed Roushdy Elkordy, Yahya H. Ezzeldin, Qingfeng Liu, Kee-Bong Song, Mostafa El-Khamy, Salman Avestimehr

arXiv:2308.06522v1cs.LGcs.AI

TL;DR

Federated fine-tuning must contend with limited client resources, while PEFT’s performance gap with full fine-tuning grows as client data becomes more diverse. The paper proposes data-driven priming for LoRA, with SLoRA as an efficient variant, and reports performance comparable to full fine-tuning at minimal communication, time, and storage cost.

  • Problem

    Fine-tuning language models in FL is constrained by edge clients’ limited resources, and the paper reports that the PEFT–full-fine-tuning gap widens as client data becomes more heterogeneous.

    Source: §1 · p. 2
  • Method

    Primed-LoRA uses data-driven priming in a two-stage PEFT approach, with an efficient SLoRA variant.

    Source: §1 · p. 2, §4 · p. 4, §4.1 · p. 4
  • Results

    SLoRA maintains performance comparable to full fine-tuning with minimal communication, time, and storage cost.

    Source: §7 · pp. 8–9
  • Takeaways & Limitations

    The paper presents SLoRA as a way to narrow the performance gap between PEFT and full fine-tuning while retaining parameter efficiency.

    Source: §1 · p. 2

    The study assumes clients have similar resource limitations and power, leaving resource heterogeneity for future work.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Transfer learning via fine-tuning pre-trained transformer models has gained significant success in delivering state-of-the-art results across various NLP tasks. In the absence of centralized data, Federated Learning (FL) can benefit from distributed and private data of the FL edge clients for fine-tuning. However, due to the limited communication, computation, and storage capabilities of edge devices and the huge sizes of popular transformer models, efficient fine-tuning is crucial to make federated training feasible. This work explores the opportunities and challenges associated with applying parameter efficient fine-tuning (PEFT) methods in different FL settings for language tasks. Specifically, our investigation reveals that as the data across users becomes more diverse, the gap between fully fine-tuning the model and employing PEFT methods widens. To bridge this performance gap, we propose a method called SLoRA, which overcomes the key limitations of LoRA in high heterogeneous data scenarios through a novel data-driven initialization technique. Our experimental results demonstrate that SLoRA achieves performance comparable to full fine-tuning, with significant sparse updates with approximately $\sim 1\%$ density while reducing training time by up to $90\%$.

1 Introduction

Fine-tuning language models in federated settings must work within clients’ limited resources while preserving data privacy. The paper examines how client-data heterogeneity affects PEFT and proposes Primed-LoRA, with an efficient SLoRA variant, to narrow its performance gap with full fine-tuning.

1 Introduction

  • Federated learning lets users train a common model locally on private data and share model updates rather than the data itself.The paper frames this as a way to support useful applications while addressing privacy concerns.Source: §1 · p. 1
  • Fine-tuning large language models can strain edge clients’ computation, bandwidth, energy, and memory, especially when multiple tasks require separate fine-tuned models.The memory required for fine-tuned models grows linearly with the number of tasks.Source: §1 · pp. 1–2
  • PEFT updates only a small set of task-specific parameters while freezing most pretrained weights, reducing the parameter budget in centralized training.The paper notes that a successful federated PEFT design could also reduce communication costs.Source: §1 · p. 2
  • As client data becomes more heterogeneous, the performance gap between PEFT and full fine-tuning widens; the paper proposes Primed-LoRA and its efficient SLoRA variant to narrow it.The introduction identifies parameter efficiency and lower training and communication costs as goals of the approach.Source: Figure 1 · p. 1, §1 · p. 2

2 Related Work

Related work covers parameter-efficient methods, federated PEFT studies, and sparse training for federated learning. This paper focuses on language models and client-data heterogeneity, and proposes an approach for extreme non-IID settings.

2 Related Work

  • PEFT methods either tune subsets of existing parameters or add task-specific modules while freezing the pretrained model.Module designs include serial bottleneck adapters, parallel LoRA modules, and prefix or prompt tuning at attention heads or embeddings.Source: §2 · p. 2
  • Prior FL studies examined PEFT for vision tasks and reported that it could replace full fine-tuning without compromising performance while reducing communication costs (Sun et al., 2022; Zhang et al., 2022).Those studies considered client stability, data distribution, and differential privacy settings.Source: §2 · p. 2
  • This study examines PEFT for language models and the effect of client-data heterogeneity on NLP performance, extending beyond benchmarking to an approach for extreme non-IID settings.The related work distinguishes it from prior work focused on vision and vision-language models.Source: §2 · p. 2
  • Federated sparse training updates either a full-size server model from sparse client training or a model kept sparse at both client and server; PEFT also addresses storage across multiple tasks.The passage describes these as distinct approaches that share a goal of reducing client training complexity.Source: §2 · pp. 2–3
  • PEFT applies very sparse task-specific updates while retaining the full pretrained model, supporting storage savings and lower communication costs with performance comparable to full fine-tuning (Hu et al., 2021; Houlsby et al., 2019; Zaken et al., 2021; Pfeiffer et al., 2022; He et al., 2021).The passage contrasts this with sparse-training methods that can require storage for task-specific sparsification patterns (Frantar and Alistarh, 2023; Gordon et al., 2020).Source: §2 · p. 3

3 Preliminaries

The preliminaries describe PEFT baselines, show that client-data heterogeneity widens the PEFT–full-fine-tuning gap, and introduce LoRA’s low-rank update design.

3.1 PEFT Baselines in Centralized Learning

  • The centralized PEFT baselines include bottleneck-module methods placed at different Transformer components, parallel LoRA modules, and BitFit, which tunes only bias terms.Pfeiffer places its module after the feed-forward block; Houlsby places it after both multi-head attention and feed-forward blocks.Source: §3.1 · p. 3

3.2 Observation: PEFT is challenged when data distribution gets non-IID

  • With greater client-data heterogeneity, the performance gap between full fine-tuning and PEFT grows, so naive local use of PEFT in FL can cause substantial performance loss.The paper reports this pattern across different models and datasets.Source: §3.2 · p. 3

3.3 Low-Rank Adaption: LoRA

  • LoRA freezes the pretrained weight matrix and trains a low-rank factorization of its update; its parallel module uses a down-projection matrix B followed by an up-projection matrix A.The module is added in parallel to the original pretrained weight matrix.Source: §3.3 · p. 3
  • LoRA can track full fine-tuning in more homogeneous federated settings, but in highly non-IID data it may fall short of its performance or converge more slowly.The cited comparison is for 20News group data and different data distributions.Source: Figure 3 · p. 4
  • LoRA scales its module output using β, which the passage describes as a constant in r that can reduce hyperparameter retuning when r changes.The module output is added coordinate-wise to the original model output.Source: §3.3 · p. 4

4 Our Proposed approach (Primed-LoRA)

LoRA can lag full fine-tuning in highly heterogeneous federated data, motivating Primed-LoRA’s two-stage approach: clients first learn an initialization, then fine-tune LoRA blocks. FLoRA uses full fine-tuning to prime those blocks, while SLoRA sparsifies Stage 1 updates to reduce their cost.

4 Our Proposed approach (Primed-LoRA)

  • In highly non-IID settings, LoRA may lag full fine-tuning or converge more slowly, motivating Primed-LoRA’s learned initialization for Stage 2.Clients collaborate in Stage 1 to find a starting point, then run LoRA with the learned initializers in Stage 2.Source: §4 · p. 4

4.1 Full fine-tuning for priming LoRA

  • FLoRA fine-tunes the full model for a few Stage 1 rounds, then uses SVD to factor the accumulated weight change into LoRA initializers.The resulting A and B matrices form a low-rank approximation of the parameter change.Source: §4.1 · p. 4
  • FLoRA’s Stage 1 has full fine-tuning’s communication and computation cost, although its Stage 2 updates only LoRA blocks.FLoRA can improve the global model, especially with more Stage 1 training rounds.Source: §4.1 · p. 4
  • Stage 1 model performance directly affects the final model, so reducing its round count trades lower cost against the quality of the final result.Figure 5 examines the impact of Stage 1 federated rounds on final Stage 2 performance for 20News group on Albert.Source: §4.1 · p. 5, Figure 5 · p. 5

4.2 Sparse Fine-tuning for priming LoRA

  • SLoRA reduces Stage 1 update and processing costs by letting clients update only weights selected by a shared server-generated mask.A common mask avoids density growth from differing client masks; its density is set above the number of parameters in the LoRA blocks.Source: §4.1 · p. 5, §4.2 · pp. 5–6
  • After sparse Stage 1, SLoRA uses SVD to decompose the accumulated update into A and B for Stage 2 LoRA.The server initializes the mask, and clients update only parameters within it.Source: Figure 4 · p. 5, §4.2 · p. 6

5 Experiments

The experiments evaluate PEFT methods for federated language-model classification across different models, datasets, and client-data distributions. They report final global-model accuracy alongside training and communication costs.

5.1 Setting

  • Experiments use Albert and DistilBERT on 20News group and News Category, with 100 clients and 20 or 10 participants per round depending on the setting.The 20News group dataset has about 19K examples and a 60/40 train/test split; News Category has 15 labels, with 10% of its roughly 330K training examples sampled.Source: §5.1 · p. 6

5.2 Metrics

  • Final global-model accuracy is the main metric; training time and communication cost are also reported.Experiments use one NVIDIA A100 GPU and average the three best results from five seeds; FedAvg aggregates client updates.Source: §5.2 · p. 6

5.3 Baselines

  • Baselines include LoRA, Houlsby, Pfeiffer, and BitFit, with data heterogeneity varied using LDA allocation and pathological non-IID data.Smaller r values yield smaller updates; BitFit's update size is fixed, and the News Category figure covers PEFT performance across settings and data distributions.Source: §5.3 · pp. 6–7, Figure 6 · p. 7

6 Evaluation

The evaluation finds that PEFT falls further behind full fine-tuning as client data becomes more heterogeneous, while SLoRA uses an SFT-derived initialization before LoRA. SLoRA achieves accuracy comparable to full fine-tuning and performs favorably against PEFT baselines in the reported settings.

6.1 Data Heterogeneity

  • As client data became more non-IID, the PEFT-to-FFT performance gap grew in all four model-and-dataset settings.The comparison covers DistilBERT and Albert on 20News group and News Category.Source: §6.1 · p. 7

6.2 Update Size

  • Larger update sizes improved final global-model accuracy on Albert for 20News group, but also increased communication cost.Table 2 reports the impact of update density for this dataset and model.Source: §6.2 · p. 7, Table 2 · p. 7

6.3 Performance of SLoRA

  • SLoRA uses two stages: clients first update the model with SFT, then use that update to initialize LoRA.The evaluation compares SLoRA with Stage 1 alone (SFT) and Stage 2 alone (LoRA).Source: §6.3 · p. 7
  • On Albert and DistilBERT for 20News group, SLoRA converged to better accuracy in fewer training rounds.Figure 7 presents the SLoRA performance for those models and dataset.Source: §6.3 · p. 8, Figure 7 · p. 8
  • SLoRA achieved accuracy comparable to full fine-tuning, with a slight marginal improvement; LoRA and SFT maximum accuracy fell −1.67 and −0.57, respectively.SLoRA also showed higher stability across seeds than the other baselines.Source: §6.3 · p. 8, Table 4 · p. 8
  • With the number of rounds fixed, SLoRA's performance gap over LoRA and SFT increased to 4% and 14.2%, respectively.The paper notes that round count matters for user availability and computational resource consumption.Source: §6.3 · p. 8

7 Conclusion

The conclusion reports that PEFT methods perform worse than full fine-tuning as client data diversity increases, while SLoRA maintains the same performance as full fine-tuning with minimal communication, time, and storage cost.

7 Conclusion

  • As data diversity increased, PEFT methods performed worse than full fine-tuning, while SLoRA maintained the same performance at minimal communication, time, and storage cost.The conclusion presents SLoRA as the approach proposed to overcome PEFT's performance limitation.Source: §7 · pp. 8–9

Limitations

The experiments assume clients have similar resource limitations and power.

Limitations

  • The evaluation assumes similar client resources and power, leaving resource heterogeneity for future work.The authors identify users with differing power and update-size needs for training or inference as a setting to consider next.Source: Limitations · p. 9

Ethics Statement

The section states that the work raises no ethical concern and describes a low-rank SVD approximation of the fine-tuning update, trading parameter savings against approximation accuracy.

Ethics Statement

A Priming LODA from FFT using SVD

B Performance and cost of 20News group on DistilBERT

For 20News on DistilBERT, SLoRA is reported to outperform FFT while retaining better training time; the section also compares its stages and reports the same observations for News category data.

B Performance and cost of 20News group on DistilBERT

B.1 Impact of Update size.

Across the reported settings, denser updates perform better, while greater data heterogeneity is associated with a considerable performance drop.

B.1 Impact of Update size.

Loading 2308.06522v1…