Source-linked AI summary
Heterogeneous LoRA for Federated Fine-tuning of On-Device Foundation Models
Yae Jee Cho, Luyang Liu, Zheng Xu, Aldi Fahrezi, Gauri Joshi
TL;DR
Federated fine-tuning of on-device foundation models must handle heterogeneous client systems and data while preserving the benefits of parameter-efficient adaptation. The paper proposes HETLORA, which assigns heterogeneous LoRA ranks and combines local rank self-pruning with server-side sparsity-weighted aggregation. HETLORA improves training speed, communication and computation efficiency, and final performance over homogeneous LoRA, while its rank-distribution assumption leaves correlated rank and data distributions for future study.
Problem
Federated fine-tuning of on-device foundation models must accommodate heterogeneous client system capabilities and data distributions, but homogeneous LoRA ranks do not fit this setting well.
Method
HETLORA assigns different LoRA ranks to clients, applies local rank self-pruning, and aggregates heterogeneous modules with sparsity-weighted aggregation.
Results
HETLORA achieves better training speed, communication and computation efficiency, and final performance than homogeneous LoRA.
Takeaways & Limitations
Heterogeneous LoRA provides a practical federated fine-tuning approach for on-device foundation models with varying client resources and data complexity.
Takeaways & Limitations
The method assumes rank distribution across clients is independent of data distribution, although these distributions may be correlated.
Abstract
from arXiv · showhide
Foundation models (FMs) adapt well to specific domains or tasks with fine-tuning, and federated learning (FL) enables the potential for privacy-preserving fine-tuning of the FMs with on-device local data. For federated fine-tuning of FMs, we consider the FMs with small to medium parameter sizes of single digit billion at maximum, referred to as on-device FMs (ODFMs) that can be deployed on devices for inference but can only be fine-tuned with parameter efficient methods. In our work, we tackle the data and system heterogeneity problem of federated fine-tuning of ODFMs by proposing a novel method using heterogeneous low-rank approximations (LoRAs), namely HetLoRA. First, we show that the naive approach of using homogeneous LoRA ranks across devices face a trade-off between overfitting and slow convergence, and thus propose HetLoRA, which allows heterogeneous ranks across client devices and efficiently aggregates and distributes these heterogeneous LoRA modules. By applying rank self-pruning locally and sparsity-weighted aggregation at the server, HetLoRA combines the advantages of high and low-rank LoRAs, which achieves improved convergence speed and final performance compared to homogeneous LoRA. Furthermore, HetLoRA offers enhanced computation efficiency compared to full fine-tuning, making it suitable for federated fine-tuning across heterogeneous devices.
1 Introduction
The paper studies federated fine-tuning of on-device foundation models under heterogeneous client resources and data, proposing HETLORA to assign different LoRA ranks across clients. HETLORA combines rank self-pruning and sparsity-weighted aggregation to improve efficiency and performance over homogeneous-rank LoRA.
- Motivation: On-device foundation models require federated fine-tuning because their smaller parameter sizes weaken downstream performance, while private client data cannot always be centralized.ODFMs are defined as models with at most a few billion parameters that fit on limited-capacity clients.
- Motivation: Existing federated parameter-efficient methods do not adapt PEFT configurations to heterogeneous client system capabilities and data distributions.The paper identifies this as a practical gap for ODFM fine-tuning.
- HETLORA: HETLORA assigns heterogeneous ranks to clients and uses rank self-pruning plus sparsity-weighted aggregation to accommodate differing capabilities and data complexities.The pretrained weights remain on-device while the server aggregates heterogeneous LoRA modules into a global module.
- LoRA trade-off: Homogeneous LoRA ranks create a trade-off: larger ranks accelerate training but overfit faster, whereas smaller ranks train more slowly without the same overfitting.LoRA reduces trainable parameters to at most 0.02% of the original ODFM size.
- Results: HETLORA outperforms homogeneous-rank LoRA in training speed, communication and computation efficiency, and final performance.The evaluation covers PaLM 2 XXS and XS on multi-session chat and Reddit summarization data from clients.
2 Related Work
Prior work combines parameter-efficient fine-tuning methods with federated learning, but limited work adapts LoRA training itself to heterogeneous device capabilities. HETLORA addresses this setting by allowing client-specific LoRA ranks.
- Parameter-Efficient Fine Tuning: Parameter-efficient fine-tuning freezes most foundation-model parameters while tuning a subset or adding trainable components such as adapters, prompts, prefixes, or low-rank approximations.
- Federated Fine-Tuning: Recent federated foundation-model work collaboratively trains prompts or adapters, but standard federated settings may not fit smaller on-device models or heterogeneous client capacities.
- LoRA in Federated Learning: Prior LoRA-for-federated-learning work studies initialization or retains standard LoRA training, without adapting the training process to heterogeneous device capabilities.
3 Federated Fine-Tuning with LoRA
Federated LoRA fine-tuning freezes pretrained ODFM weights while clients train and communicate low-rank modules. HETLORA supports heterogeneous client ranks through truncation, rank self-pruning, and sparsity-weighted aggregation.
- LoRA Federated Fine-Tuning: LoRA represents the trainable update as ∆W = BA, reducing trainable parameters while keeping the pretrained ODFM weights fixed.The factors have rank r, with B ∈ R^(d×r) and A ∈ R^(r×l).
- Heterogeneous LoRA: HETLORA assigns each client a rank within [rmin, rmax] to accommodate heterogeneous system capabilities and data complexities.Client ranks may change across communication rounds through self-pruning.
- Distribution via Truncation: The server distributes global modules by truncating their rank dimension to each selected client’s rank, then clients train locally and return updated modules.This replaces a single fixed rank with client-specific modules while preserving the federated training cycle.
- Local Training with Rank Self-Pruning: Clients self-prune ranks using a norm-based regularizer that penalizes the last ranks, with γ < 1 controlling pruning intensity.After local training, clients compare updated and initially received module norms and prune the last ranks when the updated norm is smaller.
- Sparsity-Weighted Aggregation: Sparsity-weighted aggregation reduces the bias of zero-padding toward high-rank clients by weighting reconstructed updates according to their singular-value norms.Naïve zero-padding can overemphasize noisy updates from high-rank clients; HETLORA instead de-emphasizes less informative large-rank updates.
4 Experiments
Experiments evaluate PaLM 2 ODFMs on multi-session chat and Reddit summarization, comparing homogeneous and heterogeneous LoRA with full fine-tuning. HETLORA combines heterogeneous ranks, self-rank pruning, and sparsity-weighted aggregation to improve training speed, final performance, and efficiency relative to homogeneous LoRA.
- Experimental setup: Experiments use PaLM 2 XXS and XS models on multi-session chat and Reddit summarization tasks, with perplexity and RougeL as evaluation metrics.The LoRA modules are applied to self-attention layers, and clients are sampled per communication round for each task.
- Homogeneous LoRA: Higher homogeneous LoRA ranks reach better perplexity with fewer communication rounds but overfit faster, whereas rank r = 1 trains more slowly without overfitting.This creates a trade-off between rapid performance achievement and overfitting.
- Heterogeneous LoRA: Larger heterogeneous minimum ranks speed training but worsen performance, while smaller ranks reduce overfitting, including when rmax = 50.The results attribute this reduced overfitting to the regularizing effect of smaller-rank LoRA modules.
- HETLORA: Adding self-rank pruning and sparsity-weighted aggregation prevents overfitting and improves training speed and final performance over other baselines.With rmin = 5, the combined method achieves these improvements despite the higher minimum rank.
- Comparison with baselines: HETLORA with rmin = 5 and rmax = 50 trains faster and performs better than homogeneous ranks r ∈ {5, 50} and reconstruction+SVD, while slightly underperforming full fine-tuning.It also requires significantly fewer trained and communicated parameters than full fine-tuning to reach the targeted Reddit and multi-session chat performance.
- Ablation: Pruning with γ = 0.99 outperforms no pruning at γ = 1, whereas aggressive pruning at γ = 0.85 harms performance.The reported explanation is that moderate pruning reduces noise from large-rank clients whose data complexity is low.
5 Discussions and Concluding Remarks
HETLORA addresses system and data heterogeneity in federated fine-tuning of on-device foundation models. It improves training speed, efficiency, and final performance over homogeneous LoRA, while leaving rank assignment and theoretical analysis for future work.
- HETLORA caters to device system and data heterogeneity in federated fine-tuning of on-device foundation models.
- HETLORA achieves better training speed, communication/computation efficiency, and final performance than homogeneous LoRA.
- Increasing ranks does not always improve homogeneous LoRA, consistent with prior findings.
- Effective rank assignment across clients for better convergence and performance remains an open question.
- Theoretical convergence and generalization of heterogeneous LoRA remain important next steps.
6 Limitations
The method assumes that rank distribution across clients is independent of data distribution, although these distributions may be correlated in practice.
- HETLORA assumes rank distribution across clients is independent of data distribution.
- Future work should examine correlated rank and data distributions to understand their implications for heterogeneous LoRA.