[Paper Review] Federated Fine-tuning of Large Language Models under Heterogeneous Tasks and Client Resources

Jiamu Bai, Daoyuan Chen+3 more|arXiv (Cornell University)|Feb 18, 2024

Topic Modeling4 citations

TL;DR

This paper proposes FlexLoRA, a novel federated fine-tuning framework for large language models that dynamically adjusts local LoRA ranks based on client resources, enabling full utilization of heterogeneous client capabilities. By aggregating full-size LoRA weights via SVD-based redistribution, FlexLoRA enhances global model generalization, achieving up to 3.1% average improvement in downstream NLP tasks across 1,600+ clients with diverse tasks and resources.

ABSTRACT

Federated Learning (FL) has recently been applied to the parameter-efficient fine-tuning of Large Language Models (LLMs). While promising, it raises significant challenges due to the heterogeneous resources and data distributions of clients. This study introduces FlexLoRA, a simple yet effective aggregation scheme for LLM fine-tuning, which mitigates the ``bucket effect'' in traditional FL that restricts the potential of clients with ample resources by tying them to the capabilities of the least-resourced participants. FlexLoRA allows for dynamic adjustment of local LoRA ranks, fostering the development of a global model imbued with broader, less task-specific knowledge. By synthesizing a full-size LoRA weight from individual client contributions and employing Singular Value Decomposition (SVD) for weight redistribution, FlexLoRA fully leverages heterogeneous client resources. Involving thousands of clients performing heterogeneous NLP tasks and client resources, our experiments validate the efficacy of FlexLoRA, with the federated global model achieving consistently better improvement over SOTA FL methods in downstream NLP task performance across various heterogeneous distributions. FlexLoRA's practicality is further underscored by our theoretical analysis and its seamless integration with existing LoRA-based FL methods, offering a path toward cross-device, privacy-preserving federated tuning for LLMs.

Motivation & Objective

• To address the 'bucket effect' in federated learning, where low-resource clients constrain the global model's performance despite high-resource clients having unused capacity.
• To improve the generalization of federated LLMs by enabling clients with abundant resources to contribute higher-rank LoRA adapters, reducing task-specific overfitting.
• To design a simple, plug-and-play aggregation scheme compatible with existing LoRA-based FL methods that fully leverages heterogeneous client resources.
• To provide theoretical justification and empirical validation for the effectiveness of dynamic rank adjustment and SVD-based weight redistribution in improving global model performance.

Proposed method

• FlexLoRA dynamically assigns local LoRA ranks to clients based on their available computational and memory resources, allowing high-resource clients to use larger ranks.
• It constructs a full-size LoRA weight matrix from individual client contributions before aggregation, avoiding direct parameter averaging.
• Singular Value Decomposition (SVD) is applied to redistribute the global LoRA weights, enabling effective merging of diverse rank contributions.
• The method integrates seamlessly with existing LoRA-based FL frameworks such as FedAvg and FedIT, enhancing them without architectural changes.
• A theoretical analysis supports the generalization benefits of the approach, particularly under heterogeneous data and resource distributions.
• The framework is evaluated in a large-scale, cross-device FL simulation with over 1,600 clients performing diverse NLP tasks.

Experimental results

Research questions

• RQ1 Can dynamic adjustment of local LoRA ranks mitigate the 'bucket effect' in federated LLM fine-tuning under heterogeneous client resources?
• RQ2 To what extent can SVD-based redistribution of full-size LoRA weights improve global model generalization compared to standard FedAvg or rank-constrained aggregation?
• RQ3 How does FlexLoRA perform across diverse NLP tasks with varying degrees of task-specificity and data heterogeneity?
• RQ4 Can FlexLoRA be effectively integrated into existing LoRA-based FL baselines without compromising convergence or requiring retraining?
• RQ5 What is the impact of higher-rank local LoRA updates on zero-shot generalization and downstream task performance in federated settings?

Key findings

• FlexLoRA achieved up to a 3.1% average improvement in downstream NLP task performance compared to standard FedAvg and FedIT baselines.
• The method demonstrated a 4% improvement in zero-shot Rouge-L scores and up to 4% gains in overlap extraction and textual entailment tasks.
• In tasks requiring reasoning over sentence relationships, FlexLoRA improved performance by up to 4% and 2.5% respectively, highlighting its strength in generalization.
• SVD-based approximation achieved a low error ratio of 0.16 when using only the top 30 singular values, confirming the fidelity of the redistributed weights.
• The performance of FlexLoRA closely matched that of FedIT with homogeneous high-rank settings, validating its effectiveness under controlled conditions.
• Empirical results confirmed that FlexLoRA maintains performance parity with standard methods under homogeneous conditions while significantly outperforming them under heterogeneity.